Scheduling Recursive Reasoning in Looped Transformers

Scheduling Recursive Reasoning in Looped Transformers. By Boyuan Wang and five others, from the Southern University of Science and Technology, the Chinese University of Hong Kong, Shenzhen, and one other institution. Published on the twenty-ninth of September. No code repository is linked to it.

Recurrent reasoning models have drawn growing attention as a way to scale test-time computation. Typically they refine a latent state over and over with shared parameters. Each learned update, however, is applied at a fixed unit scale. That can be too conservative when the updates are making persistent progress, and too aggressive when they fluctuate. Either way, it limits what additional loops can buy. A walker who takes the same stride whether the path is straight or zigzagging will either crawl along the straight stretches or overshoot on the turns.

To work out how the scale ought to vary along the trajectory, the authors first analyze how sensitive the terminal loss is to the scale of the recurrent update. The temporal average of that sensitivity, they show, admits an exact decomposition into two contributions. One comes from persistent progress. The other comes from centered fluctuation. The Trajectory Adaptive Progress-Fluctuation Scheduler, TAPS, is built on that split. It tracks the balance between the two across recurrent updates, and adapts the step size online. When progress dominates, the analysis points toward larger steps. When fluctuation dominates, toward smaller ones. The scheduler reads that balance as the loop runs, rather than fixing a scale in advance.

The theory comes with guarantees of a specific kind. The authors establish sufficient conditions under which TAPS reduces expected terminal loss, and reaches a target quality in fewer recurrent loops. In practice the scheduler can enter at two points. At inference, it works on a trained model without retraining. And the same progress-fluctuation principle can be incorporated into training itself, so the model learns with the same adaptive scaling it will later run under.

Applied without retraining, TAPS improves terminal accuracy across structured reasoning tasks. With the principle built into training, it yields further accuracy gains, and at matched baseline accuracy it delivers up to a one point five six times wall-clock speedup. The authors support its broad applicability with results across diverse recurrent architectures and inference strategies.

Architecture and depth are the familiar controls on recurrent inference. The authors argue for a third that complements them: update scale. Under that view, the question is not only how many loops run, but how large a step each loop takes, set by whether recent updates have been making steady progress or fluctuating.