Improving Test-Time Scaling with Adaptive Looped Transformers
Yichen You, Tianyu Fu, Aosong Feng, Xingtai Lv, Xuefei Ning, Ning Ding, Yu Wang
cs.CL, cs.LG
2026-09-29
TaH2 post-trains Qwen3 with a token-wise loop decider. On 1.7B AIME the slope rises from 1.79 to 2.74 (+53%) and beats the non-looped peak by 3.4 points at matched FLOPs.
Looped transformers reuse one stack on every token so extra depth does not grow the parameter count. Ouro loops the full stack. Huginn loops a middle block. Earlier papers match parameter counts or per-token FLOPs. They do not ask whether looping still helps once generations get longer and you plot accuracy against decoding FLOPs.
The Tsinghua and Yale group post-trains the same Qwen3-Base checkpoint into a non-looped Standard model and into looped Ouro and Huginn copies, then sweeps output cutoffs from 4K to 16K. The slope is accuracy points per doubling of decode FLOPs. Fixed-depth Ouro (M=2), gated Ouro, and Huginn (M=3) land at 2.12, 2.27 and 2.26 against Standard's 1.79. Steeper slopes, worse accuracy across the overlapping compute band. On the validation set, 52.3% of Ouro tokens move by at most 0.001 from iteration 1 to 2, and 21.7% get worse; Huginn's shares are 33.5% and 15.9%. Fixed depth bills every token for extra passes that many positions do not use.
TaH2 jointly post-trains the backbone and a token-level iteration decider. A learned updater reinjects the token embedding into the next pass. Updater plus decider add under 3% parameters. After each pass the decider reads the embedding, hidden state and TopK probabilities and emits a continue probability. The token stops below 0.5, otherwise it runs up to ceiling M.
Depth labels are online. The live backbone takes one more pass; tokens that already stopped get a no-gradient lookahead. The label is whether next-token loss actually dropped. Positive gains are ranked until they cover 99% of total positive gain; the tail is dropped as noise. Labels track the current backbone, unlike TaH's staged offline mismatch labels and unlike Ouro's freeze-then-train exit gate. Training follows the same depth policy used at inference, so the model is not trained at full depth and then asked to exit early.
The output mixes per-iteration distributions with stopping weights. The joint loss is next-token cross-entropy plus a cost-sensitive binary term on the decider (coefficient 0.05), with larger weights farther from the gain cutoff. Switching those weights to uniform drops mean AIME24-26 by 3.6 points, the largest ablation hit. Replacing the learned updater with a Top-100 embedding mix costs 2.1 points; using only the last iteration costs 1.0.
All 1.7B runs share 273K distilled prompts and a 16K context.
| Method | 10-bench avg | FLOPs/token |
| Standard (M=1) | 37.7 | 1.00× |
| Huginn (M=3 / 7) | 38.2 / 38.4 | 1.90× / 3.72× |
| Ouro (M=2 / 4) | 37.3 / 36.4 | 1.91× / 3.75× |
| TaH2-fixed (M=2 / 4) | 39.4 / 40.3 | 1.99× / 3.97× |
| TaH2 (M=2 / 4 / 8) | 40.6 / 40.9 / 42.5 | 1.21× / 1.47× / 2.37× |
On AIME24-26, TaH2 at M=2 has slope 2.74 versus Standard's 1.79, a 53% increase. With a 32K eval cutoff, Standard plateaus at 12.0% with 191.7 TFLOPs per response; TaH2 reaches 15.4% at that compute, 3.4 points higher. Raising the ceiling from 2 to 8 grows the AIME gap versus Standard from +2.8 to +3.9. Ouro and Huginn flatten. Mean realized depth at M=2 is only 1.23, which is why per-token FLOPs stay at 1.21×.
HumanEval moves from 63.6 to 72.6 at M=8; OlympiadBench from 40.4 to 47.4. At 4B and 8B, TaH2 (M=2) adds 3.2 and 2.4 points on the reported averages. The largest single jump is AIME26 at 4B, 45.5 to 52.4. AIME cons@32 sits at 27.3% to 29.3% against Standard's 21.9%. In an extended MiniSGLang server, M=2 costs 22% more FLOPs per token and 30% to 34% more end-to-end latency, with throughput near 0.71×. The accuracy-latency curve still sits above Standard.
GPQA barely moves: 31.5 for Standard, 32.0 for TaH2 at M=8, while Ouro at M=2 reaches 34.0. Gains concentrate on math and code.
This is a post-training recipe for adding recurrent depth to an existing dense LLM, not a from-scratch looped pretrain. M=2 is the practical operating point: most tokens run a little more than one pass, extra latency is about a third, and the AIME slope already pulls away. M=8 still adds accuracy; training FLOPs reach 5.17× Standard.
Naive fixed-depth looping in this post-training setup is often both more expensive and less accurate than not looping. The waste is per token, so depth has to be allocated per token. The result is still an SFT increment. Average gains shrink from 3.2 at 4B to 2.4 at 8B, and nothing larger is tested.
Training costs more than plain SFT. At 1.7B, M=2 is 1.87× Standard, and lookahead labels take 20% to 24% of training FLOPs; M=8 reaches 5.17×. The method is SFT-only. On-policy distillation and RL are left open.
The Ouro and Huginn numbers are reimplementations for this post-training setup, not the original pretrained checkpoints. The claim covers retrofitting recurrence onto a dense base, not native looped pretraining. Headline AIME numbers are avg@32 on 30-problem sets. The exit threshold is fixed at 0.5 with no serving Pareto sweep. Adaptive depth also needs a custom engine that batches mixed iteration depths in one forward pass; a vanilla vLLM path is not reported.