Improving Test-Time Scaling with Adaptive Looped Transformers

Improving Test-Time Scaling with Adaptive Looped Transformers. By Yichen You and six others, from Tsinghua University and Yale University. Published on the twenty-eighth of September. Its code repository is public and real: about twenty-five source files.

Looped transformers reuse their layers for latent computation, and they have shown promising parameter efficiency. Prior comparisons between looped and ordinary models hold parameters or per-token FLOPs fixed. What has gone largely unexamined is whether looping helps test-time scaling, as outputs grow longer. To study that, the authors post-train looped transformers and measure an accuracy-compute slope: how much accuracy is gained for each doubling of the decoding FLOPs spent at test time.

The picture that emerges is mixed. Existing looped transformers often have a steeper slope than their non-looped baseline, yet at matched compute they still perform worse. Their accuracy rises faster with compute, but from a lower place. The analysis points to why. Fixed-depth looping spends the extra iterations on every token, and many tokens get no benefit from them.

TaH2 follows directly from that observation. It lets the model focus its extra iterations on the tokens that actually benefit from looping, rather than spreading them evenly.

Deciding which tokens those are is learned. The backbone and an iteration decider are post-trained together through what the authors call lookahead depth supervision. It relies on online labels that say whether iterating further improves the prediction. The decider is thus trained on the very question it has to answer at test time: for this token, at this depth, would one more pass improve the prediction. Training the backbone jointly means the model's layers and its sense of when to reuse them develop together, rather than one being fitted around the other after the fact.

The result improves both the efficiency of test-time scaling and the accuracy it can reach, which were exactly the two things existing looped models could not deliver together. On challenging AIME benchmarks, TaH2 raises the accuracy-compute slope by fifty-three percent over the non-looped baseline, from one point seven nine to two point seven four. At matched test-time compute, it exceeds the baseline's peak accuracy by about three point four points.

Depth tells a similar story. As the maximum iteration depth increases, existing looped models largely plateau. TaH2's gain over the non-looped baseline keeps growing, from two point eight points at depth two to three point nine points at depth eight. More room to loop keeps translating into more benefit, instead of flattening out.

Looping was never short of potential. The slope was already steeper. What held it back was spending that extra depth indiscriminately. Once iterations go only where they help, the steeper slope turns into an actual lead at matched compute, and extra depth keeps paying off where it used to stall.