Improving Test-Time Scaling with Adaptive Looped Transformers

• Existing post-trained looped models can have steeper accuracy–compute slopes than Standard yet remain less accurate across the overlapping compute range.The reported slopes are 2.12, 2.27, and 2.26 for fixed-depth Ouro, adaptive Ouro, and Huginn, versus 1.79 for Standard.Source: §1 · p. 2
• TaH2 jointly post-trains the backbone and an iteration decider using lookahead depth supervision with online labels for whether further iteration improves prediction.Its depth decisions use measured gains on the current backbone during training, unlike TaH’s staged training with offline mismatch labels.Source: §1 · p. 2
• On AIME24–26, TaH2 improves the accuracy–compute slope by 53% (2.74 vs. 1.79 points per doubling) and exceeds Standard’s peak accuracy by about 3.4 points at matched decoding FLOPs.As maximum depth increases from M = 2 to M = 8, its gain over Standard grows from +2.8 to +3.9 points.Source: §1 · p. 3
• TaH2’s gains persist at 4B and 8B scales and extend beyond math to code, QA, and tool-use benchmarks.Source: §1 · p. 3

• Test-time compute can increase through longer chains of thought, repeated sampling or tree search, latent optimization, and repeated application of shared layers.Source: §2 · p. 3
• Looped-transformer studies compare recurrence under matched parameter counts, matched per-token FLOPs, or both, while architectures repeat different layer spans.Huginn loops a middle block, Ouro the full stack, and Loopies individual layers.Source: §2 · p. 3
• Adaptive-computation methods vary depth through early exit, layer skipping, or iteration counts; TaH2 learns token-dependent depth jointly with the backbone and decider using online iteration-gain labels.This differs from TaH’s offline mismatch labels and Ouro’s gate supervision with a frozen backbone.Source: §2 · p. 3

3.1 PRELIMINARIES

• The section formalizes full-stack recurrence: shared backbone Fθ updates each token’s hidden state across iterations and produces next-token probabilities, while m_t denotes that token’s executed depth.Fixed-depth models set m_t = M for every token; Standard uses m_t = 1.Source: §3 · p. 4, §3.1 · p. 4

3.2 TEST-TIME SCALING BEHAVIOUR

• The scaling comparison uses Qwen3-1.7B-Base models post-trained on the same data, with output cutoffs from 4K to 16K and mean AIME24–26 accuracy measured over 32 samples per problem.The study fits accuracy against log2 decoding FLOPs per response.Source: §3.2 · p. 4
• Fixed-depth Ouro, adaptive Ouro, and Huginn achieve slopes of 2.12, 2.27, and 2.26 accuracy points per compute doubling, versus 1.79 for Standard, yet all trail Standard over overlapping compute.Ouro’s exit gate lowers decoding cost but does not close the accuracy gap.Source: §3.2 · p. 4
• 52.3% of Ouro tokens and 33.5% of Huginn tokens change by at most 10^-3 loss; 21.7% and 15.9%, respectively, worsen by more than that threshold.Figure 3 groups tokens by whether they become worse, change little, or improve from first to final iteration.Source: §3.2 · p. 4, Figure 3 · p. 4
• Fixed-depth looping spends iterations on many tokens that gain little or worsen, while Ouro’s exit gate remains less accurate than Standard at matched compute.Because the gate’s backbone is trained only at full depth, early exit creates a train–inference mismatch; this motivates jointly learning token-dependent depth from measured loss reductions.Source: §3.2 · p. 4

4.1 ARCHITECTURE

• TaH2 combines a shared Transformer backbone, learned input-injection updater, and token-level iteration decider; the updater and decider add fewer than 3% parameters at studied scales.The updater reinjects the input embedding while mapping the previous final-layer state to the backbone input space.Source: §4.1 · p. 5
• Each iteration runs all backbone layers with extended duo-causal attention, and lookahead queries cannot attend to other tokens’ lookahead states.During training, stopped tokens take an additional no-gradient iteration for supervision.Source: §4.1 · p. 5
• Continue probabilities determine token stopping depths and induce weights over executed depths; training follows the same depth-selection rule used at inference.This differs from Ouro’s full-depth training and TaH’s oracle-policy training.Source: §4.1 · p. 5

4.2 TRAINING WITH LOOKAHEAD DEPTH SUPERVISION

• For tokens that execute iteration m, online labels use the change in prediction loss: positive gains favor continuing, while negative gains favor stopping.Stopped tokens receive a no-gradient lookahead to obtain the next-iteration loss, and later labels remain zero after a stop label.Source: §4.2 · pp. 5–6
• At each iteration, positive gains are ranked and selected until they account for coverage ρ, excluding marginal improvements that can produce noisy labels.The smallest selected gain defines the cutoff; once a token receives a stop label, subsequent labels remain zero.Source: §4.2 · p. 6
• The joint objective combines next-token prediction loss with cost-sensitive supervision of the decider.The decider is supervised on tokens that execute the iteration, with stronger supervision farther from the cutoff.Source: §4.2 · p. 6

5.2 PERFORMANCE

• Across ten benchmarks, TaH2's average gain over Standard grows from +2.9 points at M = 2 to +4.8 at M = 8, while Ouro and Huginn stay close to Standard.On validation loss, adaptive TaH2 reaches −0.0106 versus Standard at M = 8; Ouro and Huginn remain at or above Standard at every ceiling.Source: Table 1 · p. 7, §5.2 · p. 7
• At 1.7B, TaH2 gains 2.74 points per doubling of decoding FLOPs versus 1.79 for Standard; at matched compute, it reaches 15.4% versus Standard's 12.0% peak.The matched-compute comparison uses 32K-token evaluation; TaH2's advantage is 3.4 points.Source: §5.2 · p. 8
• In actual serving, TaH2 achieves better test-time scaling and higher attainable accuracy despite 22% more decoding FLOPs per token and 30–34% higher end-to-end latency than Standard.Figure 5 plots accuracy against mean end-to-end latency, with lighter segments extending beyond the 16K training length.Source: §5.2 · p. 8, Figure 5 · p. 8

5.3 ADDITIONAL SCALES

• At 4B and 8B, TaH2 outperforms Standard on every benchmark, raising average accuracy by 3.2 and 2.4 points, respectively.On AIME math, gains reach up to 6.9 points at 4B and 4.4 points at 8B.Source: §5.3 · p. 8, Table 3 · p. 8

5.4 DESIGN CHOICE EXPLORATION

• Gain-based depth labels outperform top-1 mismatch labels by 1.1 average points, while gain-magnitude weighting, filtering marginal gains, and thresholded training decisions also improve accuracy.Uniform decision weighting causes the largest drop, 3.6 points; retaining all positive gains reduces average accuracy by 1.9 points, and sampling decisions reduces it by 1.3 points.Source: §5.4 · p. 9
• Replacing the learned updater with a probability-weighted sum of top-100 token embeddings reduces average accuracy by 2.1 points; using only the final prediction reduces it by 1.0 point.These comparisons support learned state updates and the stopping-weighted mixture, respectively.Source: §5.4 · p. 9

5.5 FURTHER ANALYSIS

• Tokens with near-zero continue probability show negative or negligible second-iteration gains, while mean gain increases with continue probability overall.At the 0.5 decision threshold, the corresponding loss reduction is near zero.Source: Figure 6 · p. 9, §5.5 · p. 9
• In sampled math and code responses, expressions and final code use fewer iterations than preceding reasoning, whereas the QA example maintains greater depth.Three representative attention heads also show distinct preferences for first-iteration states, later-iteration states, or both.Source: §5.5 · p. 9

6 CONCLUSION

• TaH2 improves the AIME accuracy–compute slope by 53%, exceeds Standard's peak accuracy by about 3.4 points at matched compute, and increases its gain from 2.8 to 3.9 points as maximum depth rises from 2 to 8.The reported gains also extend to 4B and 8B models and to code, QA, and tool-use domains.Source: §6 · p. 9
• The method is studied only under SFT; extension to on-policy distillation and reinforcement learning is left for future work.TaH2 also incurs more training FLOPs than standard SFT, though post-training requires substantially less compute than pretraining.Source: §6 · p. 9
• The authors report using generative AI only to polish the paper text, and state that the study involves no human subjects or sensitive personal data.They say AI-assisted text was checked against experimental records and cited sources.Source: AI USE STATEMENT · p. 10, ETHICS STATEMENT · p. 10
• The paper reports using publicly available base models and data, and says appendices specify the loss, architecture, baselines, FLOPs accounting, training hyperparameters, and evaluation protocol.Training and evaluation code, configuration files, and checkpoints are slated for release upon publication.Source: REPRODUCIBILITY STATEMENT · p. 10, Table 5 · p. 19

A.1.2 TRAINING DATA

• The 1.7B runs use regenerated AM-Qwen3-Distilled and Nemotron-Agentic-v1 responses from Qwen3-8B; 4B and 8B runs use Qwen3-235B-A22B responses plus tool-use samples.The 4B and 8B training sets exclude prompt-response pairs exceeding 16,384 tokens; those models train on 8B and 16B tokens, respectively.Source: §A.1.2 · p. 19

A.2.1 ARCHITECTURE DEFINITIONS

• TaH2 keeps a separate KV cache per iteration and uses duo-causal attention, allowing parallel training and prefill while restricting lookahead queries from seeing other tokens’ lookahead states.The updater fuses each token’s original embedding with its current hidden state; the decider also uses the embedding and hidden state, along with the largest prediction probabilities, to decide whether to continue.Source: §A.2.1 · pp. 19–20, Figure 7 · p. 20
• Ouro loops the full backbone, while Huginn repeats a middle block and reinjects the pre-loop representation; TaH2-fixed removes the decider and runs all M iterations at every token.The appendix adapts Ouro and Huginn to the same checkpoint and training setup as TaH2, without TaH2’s updater, decider, or online supervision.Source: §A.2.1 · pp. 20–21

A.2.2 ARCHITECTURE SUMMARY AND PARAMETER COUNTS

• The updater and decider together account for less than 3% of total parameters across the reported architecture summary.Tables 6 and 7 summarize the 1.7B architectures and module parameter counts across backbone scales.Source: §A.2.2 · p. 21, Table 6 · p. 21, Table 7 · p. 21

A.3 EVALUATION PROTOCOL

• Scaling curves truncate and re-evaluate responses at specified output-token cutoffs; the 4K–16K Figure 1a curves fit accuracy against log2 decoding FLOPs with R2 values of 0.997 for Ouro, 0.994 for Huginn, and 0.955 for Standard.Validation loss and token-level analyses use 1,000 held-out training-mixture samples scored under teacher forcing.Source: §A.3 · p. 21
• Runtime measurements batch requests at different iteration depths in shared forward passes and report decoding GFLOPs per generated token, mean latency, and aggregate output-token throughput.The AIME26 evaluation uses 30 problems, a 32K output-token limit, and one A800-80GB GPU.Source: §A.3 · p. 22

A.4 COMPUTE ACCOUNTING

• Compute accounting counts model matrix-multiplication FLOPs, treating each multiply-add as two operations, while excluding elementwise operations, sampling, memory traffic, and serving overhead.A linear map from a to b features is counted as 2ab FLOPs per token.Source: §A.4.1 · p. 22, §A.4 · p. 22

A.4.1 PRELIMINARIES

• Per-call estimates include backbone, LM-head, updater, decider, and attention costs; the backbone pass dominates, while the updater and decider each cost less than 2% of it.Attention cost varies with the number of keys a query attends to.Source: §A.4.1 · p. 22, Table 9 · p. 23

A.4.2 DECODING FLOPS

• Decoding FLOPs sum per-token pass costs using each token’s executed depth and visible KV entries; fixed-depth baselines execute all M iterations at every position.The appendix distinguishes causal-attention TaH2-fixed from models that retain a KV stream per iteration, including Ouro and Huginn.Source: §A.4.2 · pp. 23–24

A.4.3 TRAINING FLOPS

• TaH2’s training-FLOP estimate separates gradient-tracked forward computation from no-gradient online-label computation, and the comparison excludes shared implementation-level choices such as activation checkpointing.Backward cost is approximated as twice the corresponding forward cost.Source: §A.4.3 · p. 24
• For a token that stops before M, online labeling runs a no-gradient next-iteration lookahead and reconstructs context by rerunning all tokens that executed the preceding iteration.The lookahead applies the updater, backbone, and LM head; it must include the next-depth KV states of earlier tokens that continue.Source: §A.4.3 · p. 24
• At M = 2, TaH2 training costs 1.87× Standard, below TaH2-fixed at 2.01× and Ouro at 2.00× the same ceiling.Online-label lookahead accounts for 20–24% of TaH2 training FLOPs.Source: §A.4.3 · p. 25, Table 10 · p. 25

B.1 TEACHER SELECTION

• The Qwen3-8B teacher produced the highest mean AIME24–26 student accuracy, exceeding Qwen3-32B and Qwen3-235B-A22B by 1.63 and 1.80 points, respectively.The comparison fine-tuned Qwen3-1.7B-Base students with identical recipes and evaluations; Qwen3-8B was selected for the main experiments.Source: §B.1 · p. 25, Table 11 · p. 25

B.2 INFERENCE DEPTH OUTSIDE THE TRAINING CEILING

• Lowering the inference ceiling to 4 reduces AIME24–26 accuracy by 2.7–3.1 points across the benchmarks.Table 12 reports accuracy and mean realized depth for inference ceilings changed without retraining.Source: §B.2 · p. 25, Table 12 · p. 26
• Raising the ceiling to 12 leaves accuracy essentially unchanged while increasing mean realized depth by 20–27%.The reported results indicate that additional inference depth beyond the trained ceiling neither breaks the model nor improves it.Source: §B.2 · p. 25, Table 12 · p. 26

C.1 VALIDATION LOSS DURING POST-TRAINING

• Both TaH2 variants maintain lower validation loss than Standard, with greater iteration depth reducing loss further and adaptive TaH2 reaching the lowest final loss.Figure 8 presents validation loss during post-training at 1.7B; negative inset differences indicate improvement over Standard.Source: Figure 8 · p. 26, §C.1 · p. 26

C.2 PER-BENCHMARK TEST-TIME SCALING

• The per-benchmark scaling figure covers six math benchmarks at nine output-token cutoffs from 4K to 32K, but the supplied text gives no benchmark-specific outcomes.The cutoffs match Figure 4b's sampling grid.Source: §C.2 · p. 26, Figure 9 · p. 27

C.4 TOKEN-LEVEL DEPTH ALLOCATION

• Across the displayed spans, mean depth falls from 2.46 to 1.29 in OlympiadBench and from 3.31 to 1.04 in HumanEval.The sampled math and code responses shift from deeper natural-language reasoning toward shallower mathematical expressions and final code; the GPQA example retains depth through its concluding explanation.Source: §C.4 · p. 27, Figure 11 · p. 28

C.3 PARALLEL TEST-TIME SCALING

• 27.3% at n = 32 for TaH2 M = 2 versus 21.9% for Standard; TaH2 stays above every fixed-depth baseline at matched FLOPs.The comparison uses majority voting over sampled responses and decoding FLOPs per problem.Source: §C.3 · p. 27, Figure 10 · p. 28

C.5 CROSS-ITERATION ATTENTION

• At query iteration 8, three representative heads prefer first-iteration states, later-iteration states, or both.The attention analysis uses 100 randomly sampled validation sequences under teacher forcing while retaining the decider’s actual routes.Source: §C.5 · p. 27, Figure 12 · p. 29

D FULL FORMULATION OF LOOKAHEAD DEPTH SUPERVISION

• The coverage cutoff retains a specified fraction of total positive gain, while bounded cost-sensitive weights preserve supervision near the cutoff and limit large loss-change influence.The coverage cutoff and class-balancing weights are batch-calculated; clipping bounds stay fixed across experiments.Source: D FULL FORMULATION OF LOOKAHEAD DEPTH… · p. 29

• Test-time scaling methods vary reasoning length or select among candidates, while latent-space approaches alter representations or add hidden sequence positions.The cited approaches include online confidence budgets, process-reward search, effort-aware scheduling, and parallel latent blocks.Source: E ADDITIONAL RELATED WORK · pp. 29, 31
• Recurrent architectures differ in what they repeat and retain, from shared transformations or repeated blocks to exposing earlier states and mixing predictions across depths.TaH2 uses input injection between iterations and mixes output distributions across executed depths, with the final depth absorbing remaining stopping mass.Source: E ADDITIONAL RELATED WORK · p. 31
• Recurrence comparisons depend on the resource held fixed; this study varies output-token budgets after post-training and measures accuracy against decoding FLOPs per response.Prior comparisons instead include fixed unique parameter count, fixed effective depth, matched per-token FLOPs, and matched wall-clock training budgets.Source: E ADDITIONAL RELATED WORK · p. 31
• Training at greater depth and running beyond the trained depth are distinct settings, and prior work reports saturation or degradation in some beyond-training-depth evaluations.The paper’s depth-scaling experiments increase maximum iteration depth during training; inference beyond that depth is examined separately.Source: E ADDITIONAL RELATED WORK · p. 31
• Adaptive-computation methods vary by width or depth, decision timing, and training signal; TaH2 derives token-level targets online from the evolving backbone during joint training.Prior looped approaches include initial-state depth prediction, per-iteration decisions, compute penalties, terminal rewards, and explicit gain supervision.Source: E ADDITIONAL RELATED WORK · p. 32
• Serving adaptive depth may require batching, cross-loop parallelism, or KV sharing, so FLOPs savings do not by themselves establish wall-clock gains.The paper reports decoding FLOPs as its primary metric and verifies throughput with a depth-batched engine.Source: E ADDITIONAL RELATED WORK · p. 32