Designing Training Objectives for Iterative Reasoning Agents: Dense Supervision as an Adaptive Mechanism
Bryan Truong ⋅ Hoyeon Chang ⋅ Alexander Rubinstein ⋅ Seong Joon Oh
Abstract
Modern AI agents increasingly rely on iterative computation - chain-of-thought, looped Transformers, and recurrent reasoners - to ``think longer'' before acting. The training objective these systems are optimized under is itself a design choice that shapes the agent's downstream behavior: which intermediate states are reachable, whether reasoning extrapolates beyond training-time horizons, and whether the agent can halt adaptively. The standard recipe - supervising only the final iterate (\emph{endpoint supervision}) - treats this design problem implicitly and, we argue, badly. Viewed as a mechanism that the training procedure imposes on the learner's internal dynamics, endpoint supervision induces two coupled pathologies: (i) gradient signals at early reasoning steps are noisy in magnitude and unreliable in direction, and (ii) intermediate states are left unconstrained, so the agent learns a fixed-horizon strategy that breaks when more reasoning is permitted at test time. We introduce Dense Intermediate Consistency for Endpoints (DICE), a training-time mechanism that attaches a \emph{shared} readout head to every iterate and applies auxiliary losses across the full trajectory. The shared readout enforces \emph{anytime decodability}: every intermediate state is a valid action under the agent's output interface, making variable-depth reasoning well-defined rather than emergent. Across three representative iterative architectures, DICE yields consistent gains - near-perfect algorithmic extrapolation on prefix sums, $+7.4$ pp exact-match accuracy on maze solving, and a $6.6\times$ inference-time speedup via adaptive halting.
Chat is not available.
Successful Page Load