PaperlayerView exact revision on arXiv

Source-linked AI summary

Think-at-Hard: Dynamic Looped Transformers for Improved Reasoning

Tianyu Fu, Yichen You, Zekai Chen, Guohao Dai, Huazhong Yang, Yu Wang

arXiv:2511.08577v4cs.CLcs.AIcs.LGcs.PF

TL;DR

Small models have limited computation per token, and looped transformers may revise correct predictions into errors. TaH selectively iterates on hard tokens using a decider, specialized architecture, and static-oracle training. Across nine benchmarks, it improves on single-iteration and always-iterate baselines while skipping iterations on most tokens.

  • Problem

    With limited computation per output token, small models can mispredict hard tokens, and errors in critical reasoning tokens can derail the reasoning trajectory.

    Source: §1 · p. 1

  • Method

    TaH combines a neural iteration decider, duo-causal attention, depth-specific LoRA, and static-oracle training to focus deeper iterations on tokens mispredicted by an SFT reference.

    Source: §1 · p. 2, §4 · p. 4, §5.3 · p. 5

  • Results

    Across nine reasoning benchmarks, TaH gains 3.8–4.4% over Ouro while reducing latent iterations by 93%, and gains 3.0% and 3.8% over Standard for 0.6B and 1.7B models.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Selective iteration can improve reasoning quality while reducing computation by applying extra latent depth to a few tokens.

    Source: §1 · p. 2

    Because of resource constraints, the work focuses on SFT with public data and leaves online distillation and reinforcement learning for future work.

    Source: §A.6 · p. 25

Abstract

from arXiv · show

Improving the reasoning abilities of Large Language Models (LLMs), especially under parameter constraints, is crucial for real-world applications. Looped transformers address this by performing multiple latent iterations to refine each token beyond a single forward pass. However, we identify a latent overthinking phenomenon: most token predictions are already correct after the first pass, but are sometimes revised into errors in later iterations. We ask whether selectively skipping latent iterations can improve accuracy, and reveal significant potential with an oracle iteration policy that boosts performance by up to 7.3%. Motivated by this, we propose Think-at-Hard (TaH), a looped transformer optimized for selective iteration. TaH employs a lightweight neural decider to trigger latent iteration, only at tokens likely to be incorrect after the standard forward pass. During latent iterations, depth-aware Low-Rank Adaptation (LoRA) modules shift the objective from general next-token prediction to focused hard-token refinement. A duo-causal attention mechanism extends attention from the token sequence dimension to an additional iteration depth dimension, enabling cross-iteration information flow with full sequential parallelism. Experiments on nine benchmarks show consistent gains across math, QA, and coding tasks. With identical parameter counts, TaH outperforms always-iterate baselines by 3.8-4.4% while skipping iterations on 93% of tokens, and exceeds single-iteration Qwen3 baselines by 3.0-3.8%. When allowing <3% more parameters from LoRA and decider, the gains further increase to 5.3-6.2% and 6.1-6.8%, respectively. Our code is available at https://github.com/thu-nics/TaH.

1. Introduction

Small models face limited computation per token, while looped transformers can also revise correct predictions into errors. TaH selectively iterates on hard tokens, with a neural decider, specialized architecture, and static-oracle training.

1. Introduction

  • Looped transformers can correct initially wrong predictions but also flip correct ones; most tokens are already correct after the first pass.Additional reasoning steps can degrade answers, motivating selective iteration on a few tokens.Source: §1 · pp. 1–2
  • An oracle that iterates only on initially mispredicted tokens achieves up to 32% higher accuracy than always iterating with an optimized architecture.The oracle results reveal potential in selectively skipping iterations.Source: §1 · p. 2
  • Selective iteration must preserve cross-depth access and sequence-level parallelism, adapt across iterations, and stabilize decisions coupled to prediction quality.The iteration policy affects prediction quality, which in turn affects the policy's decisions.Source: §1 · p. 2
  • TaH uses a neural decider, duo-causal attention, and depth-specific LoRA to select tokens for iteration and focus deeper passes on refinement.The LoRA adapters are applied only at iterations d > 1.Source: §1 · p. 2
  • A static oracle policy supports stable training by decoupling model adaptation from policy learning and addressing their circular dependency.The policy is coupled to prediction quality, which depends on the depths assigned by that policy.Source: §1 · p. 2
  • 3.8–4.4% gains over Ouro accompany 93% fewer latent iterations; TaH also improves over single-iteration baselines across nine reasoning benchmarks.For 0.6B and 1.7B models, average gains over Standard are 3.0% and 3.8%, respectively.Source: §1 · p. 2

2. Related Work

Related approaches steer computation with token-space controls, optimize latent representations, or loop over latent iterations. TaH builds on looped transformers and makes iteration selective.

2. Related Work

  • Signal-guided methods add filler or [PAUSE] tokens for extra decoding compute, but offer limited latent control.Examples include filler tokens such as dots and learnable [PAUSE] tokens.Source: §2 · p. 2
  • Latent-optimization methods reason in internal representations with little intermediate text, but sacrifice interpretability and may require mitigation from verbal LLMs.Approaches include progressive replacement, hidden-state alignment, and logit-weighted embeddings.Source: §2 · p. 2
  • Looped transformers add latent iterations before verbalization and reuse hidden states or logit-weighted embeddings, but uniform depth risks overthinking and burdens computation.Prior work focuses on scaling iteration depth.Source: §2 · p. 2
  • TaH uses existing pretrained models with depth-aware LoRA and duo-causal attention, unlike concurrent dynamic-recursion works that degrade at non-maximum depths and require full retraining.The cited concurrent works are Bae et al. (2025) and Zhu et al. (2025).Source: §2 · p. 3

3. Preliminaries

The preliminaries describe autoregressive prediction and causal attention, then introduce latent iteration before token verbalization. They also specify comparison labels for iteration strategies.

3. Preliminaries

  • Autoregressive LLMs process the input sequence in parallel during prefill, then consume one new token with cached history during decoding.The model predicts the next token at each decode step.Source: §3 · p. 3
  • Causal attention lets each position attend to itself and earlier positions, supporting parallel next-token training and efficient inference through cached past Key/Value states.The passage contrasts parallel training with token-by-token generation.Source: §3 · p. 3
  • Looped transformers iterate over latent hidden states before verbalizing each output token, updating the current iteration with causal attention.The iteration depth is denoted d.Source: §3 · p. 3
  • Table 1 compares iteration strategies; Always1 means no latent iteration, Oracle iterates on 12–19% of tokens, and MMLU100 is the first 100 MMLU-STEM questions.Source: Table 1 · p. 3

4. Selective Latent Iteration Oracles

The oracle policy tests selective iteration by extending computation only when the reference model misses a token on its first pass. Its evaluations show that skipping iterations can improve next-token and downstream performance.

Oracle iteration policy

  • The oracle iterates only when the reference model mispredicts the target token on its first pass; the evaluation uses a greedy, locally optimal policy.The experiments analyze Ouro and TaH fine-tuned from Qwen3-1.7B-Base on a balanced 100K-sample Open-R1 subset.Source: §4 · p. 3

Next-token prediction

  • 2.1% higher next-token prediction accuracy than always iterating twice, while skipping iterations on 81–88% of tokens.The reported gain comes from avoiding latent overthinking, in which a later pass can turn a correct prediction into an error.Source: §4 · p. 4

Downstream tasks

  • 2.0–7.3% improvement in downstream performance under the oracle policy.Ground-truth tokens are unavailable during generation, so Qwen3-8B top-1 predictions serve as proxy labels.Source: §4 · p. 4

TaH design objectives

  • The oracle results motivate selective depths and deeper-iteration training focused on failing tokens, although ground-truth tokens are unavailable at inference.The passage notes that other globally optimal policies may exist and proposes neural networks as a way to approximate the oracle.Source: §4 · p. 4

5. TaH Design

TaH combines selective latent iteration with cross-depth attention, depth-specific refinement, and training aligned to an oracle policy. Its design aims to preserve parallelism while concentrating deeper computation on hard tokens.

5. TaH Design

  • TaH combines duo-causal attention, depth-adaptive LoRA, and a two-stage training scheme to support selective latent iteration.The design uses a decider to choose whether tokens iterate or verbalize.Source: §5 · p. 4

5.1. Duo-Causal Attention

  • Duo-causal attention lets tokens access earlier positions and shallower depths, preserving full parallel training while supplying cross-depth context.Each depth computes representations simultaneously using information from depths 1 through that depth.Source: §5.1 · p. 4

5.2. Depth-adaptive Model Architecture

  • With over 73% of next tokens correct at iteration one, deeper passes target refinement: LoRA applies only at d > 1, while an MLP decider chooses continuation or verbalization.The decider uses shallow, middle, and final hidden states to predict continuation probability.Source: §5.2 · p. 5

5.3. Training Scheme

  • The training scheme uses a fixed oracle policy to separate backbone adaptation from iteration-policy learning, addressing their coupled training instability.The oracle triggers another iteration when the SFT model’s top-1 prediction differs from the ground-truth token; an on-policy oracle was empirically unstable.Source: §5.3 · p. 5
  • Backbone supervision runs at oracle-determined depths to preserve first-pass accuracy on easy tokens and refine hard ones, then a frozen-backbone decider imitates the oracle.Decider imitation uses weighted binary cross-entropy.Source: §5.3 · pp. 5–6

6. Experiment

Experiments evaluate TaH’s performance, efficiency, design choices, and behavior across reasoning benchmarks. They compare TaH with standard and looped-transformer baselines, then examine its training and iteration mechanisms.

6.1. Setup

  • Training uses balanced 100K-example Open-R1 data for five epochs; evaluation covers nine benchmarks under zero-shot CoT with pass@1.The setup uses Qwen3-Base backbones, filters long responses, and uses maximum sequence lengths of 8,192 tokens except for 4B models, which use 4,096; TaH’s parameter-matched setup prunes one backbone layer, while TaH+ adds less than 3% parameters.Source: §6.1 · p. 6

6.2. Performance

  • TaH averages 1.07 iterations per token, with 4-5% additional overhead over Standard versus AlwaysThink’s 2.08-2.18× FLOPs and memory access.On AIME25 at 8K length, TaH iterates twice on 6% of tokens, with 1.48× lower memory overhead and 2.48× faster decoding than AlwaysThink; during training, it converges faster than Standard with lower validation perplexity, and its decider reaches about 83% oracle-decision accuracy.Source: §6.2 · pp. 6–7, Table 3 · p. 7, Figure 4 · p. 7
  • 3.0% and 3.8% average gains over Standard for 0.6B and 1.7B; TaH+ gains are 5.3% and 6.2%, respectively.Against Ouro, TaH gains 3.8-4.4% and TaH+ gains 6.1-6.8%; 4B TaH and TaH+ gain 1.7% and 2.2%.Source: §6.2 · p. 6, Table 2 · p. 7

6.3. Design Choice Exploration

  • Decider-based iteration outperforms always-1 and always-2 alternatives; replacing duo-causal attention drops results by 5.4% or 8.5%.Removing LoRA and residual connections also causes consistent drops, supporting their beneficial roles in objective transition across iterations.Source: §6.3 · pp. 7–8, Table 4 · p. 8
  • Token-plus-latent supervision underperforms final-output-only supervision, while dynamic oracle recalculation causes training collapse; top-1 mismatch gives the best discrepancy metric.The decider-first training alternative also faces a coupling challenge, supporting the choice of a static oracle policy.Source: §6.3 · p. 8
  • Labels from a smaller 0.6B reference and labels with 10% random flips still improve over Standard-1.7B; the clean same-scale reference gives the strongest average gain.TaH is robust across continuation thresholds and iteration ratios; the experiments use cthreshold = 0.9.Source: §6.3 · p. 8, Figure 6 · p. 8, Table 5 · p. 9
  • On the 1.7B math subset, TaH-3 (Dmax=3) and TaH-4 (Dmax=4) raise average gain over Standard from TaH-2’s +5.0% to +6.0% and +6.2%, respectively, with little extra cost.Sparse allocation of deeper iterations is more effective than uniform allocation.Source: §6.3 · pp. 8–9
  • 2.1-3.1% average gains over Standard remain when 1.7B TaH models trained with Dmax = 2 are forced to stop after their first inference iteration.This restricted-depth setting performs below full TaH inference but still improves over Standard.Source: §6.3 · p. 9, Table 7 · p. 9

7. Conclusion

TaH is presented as a selective latent-iteration method for improving reasoning performance and efficiency. Across nine benchmarks, it reports accuracy gains with minimal overhead.

7. Conclusion

  • TaH combines a neural iteration decider, duo-causal attention, depth-specific LoRA, and two-stage training under a static oracle policy.The conclusion presents this as a selective latent-iteration method that improves reasoning performance and efficiency.Source: §7 · p. 9
  • Across nine benchmarks, TaH improves accuracy by 5.3-6.2% over strong baselines with minimal overhead.The conclusion frames this as a direction for improving reasoning under parameter constraints.Source: §7 · p. 9
  • The paper presents selective iteration as a way to improve reasoning performance and efficiency under parameter constraints.Source: §7 · p. 9

Impact Statement

The paper reports performance improvements for parameter-constrained language models through selective iteration. It describes this as an early exploration with potential relevance to broader access to capable AI systems.

Impact Statement

  • Selective iteration improves performance for parameter-constrained language models.Source: Impact Statement · p. 10
  • The authors describe this as an early exploration, with potential gains comparable to much larger models while skipping iterations on most tokens.Source: Impact Statement · p. 10
  • The paper suggests that improving smaller models may help democratize access to capable AI systems across diverse settings.Source: Impact Statement · p. 10

A.1.1. TRAINING RECIPE

The training recipe uses a balanced 100K-example mixture of math, code, and QA data, with separate data preparation for the main experiments and design explorations.

A.1.1. TRAINING RECIPE

  • The main performance experiments use 100K balanced examples sampled from math, code, and QA after filtering outputs longer than 8K tokens.The filtered data contains 480M tokens, with 1% reserved for validation.Source: §A.1.1 · p. 13
  • Design-choice experiments use 75K math training samples filtered by the same 8K-output-length limit.Source: §A.1.1 · p. 13
  • Maximum sequence length is 8,192 tokens for 0.6B and 1.7B models, and 4,096 for 4B Standard, TaH, and TaH+.Ouro uses a 3,072-token maximum for 4B models due to memory constraints.Source: §A.1.1 · p. 13

A.1.2. BASELINE SETUPS

The section lists training hyperparameters and describes the AlwaysThink and Ouro comparison baselines.

A.1.2. BASELINE SETUPS

  • Table 8 presents the training hyperparameters.The supplied passage does not specify their values.Source: Table 8 · p. 13
  • AlwaysThink uses TaH’s architecture but iterates twice at every token and uses standard causal attention.Its attention is limited to the current iteration depth rather than using duo-causal attention.Source: §A.1.2 · p. 13
  • Ouro follows its official architecture and is fine-tuned on Open-R1 from Qwen-Base initialization for comparison.The setup uses Ouro’s entropy-regularized loss; results may differ from the original paper because its training data and initialization differ.Source: §A.1.2 · p. 13

A.1.3. ITERATION DEPTH FLEXIBILITY

This section examines TaH’s flexible iteration depths, evaluation settings, efficiency, and learned iteration-decider behavior, alongside token-selection and cross-iteration attention patterns. It also specifies duo-causal attention and discusses TaH’s relation to layer skipping, limitations, and illustrative examples.

A.1.3. ITERATION DEPTH FLEXIBILITY

  • For Dmax = 2, tokens that miss top-1 prediction at the first iteration are assigned depth Dmax, while easy tokens verbalize at depth 1; deeper budgets use loss-ranked quantile mapping to a target depth distribution.The iteration decider receives continuation labels at each depth and uses predicted continuation probabilities with a threshold at inference.Source: §A.1.3 · pp. 13–14

A.2.1. MATH-ONLY TRAINING

  • TaH improves math-benchmark accuracy over Standard by 4.0%, 5.0%, and 3.8% for 0.6B, 1.7B, and 4B models; TaH+ gains 5.3%, 5.4%, and 4.2%.TaH gains 8.1–11.3% over AlwaysThink at 0.6B and 1.7B; TaH+ gains 8.5–12.6%. TaH+ adds less than 3% parameters; one backbone layer is removed from TaH for parameter-budget fairness.Source: §A.2.1 · p. 14, §A.1.4 · p. 14, Table 10 · p. 15

A.2.2. ADDITIONAL ORACLE ANALYSIS

  • The oracle-controlled policy increases correct next-token revisions by selectively iterating on hard tokens, while AlwaysThink produces more incorrect than correct revisions.In MATH100 generation with Qwen3-0.6B and proxy labels from DeepSeek-R1-Distill-Qwen-32B, the trained decider beats Standard and AlwaysThink, but the ideal oracle performs better still.Source: §A.2.2 · p. 15, Figure 9 · p. 16, Table 12 · p. 16

A.2.3. ADDITIONAL BASELINES

  • Under matched training and evaluation, TaH outperforms both Standard and Standard-Pruned, while pruning alone slightly degrades performance; fine-tuned PonderingPythia-1.4B substantially underperforms.The pruning comparison isolates the gain from selective latent iteration; the Ponder result may reflect the limited capability of its backbone.Source: §A.2.3 · p. 16, Table 13 · p. 16

A.2.5. THEORETICAL EFFICIENCY ANALYSIS

  • TaH incurs only 1.04–1.05× Standard’s computation and memory-access costs, versus 2.19–2.27× for AlwaysThink.The analysis attributes TaH’s overhead mainly to average iteration depth; the current implementation is not system-optimized, with Python-level control-flow overhead from LoRA switching and bookkeeping.Source: §A.2.4 · p. 17, §A.2.5 · p. 18, Table 21 · p. 21

A.3.1. TRAINING SCHEME DYNAMICS

  • Token-only supervision with a fixed oracle policy trains stably, whereas token-plus-latent supervision and decider-based training lag, and the dynamic oracle is unstable during generation.Under matched continuation budgets, top-1 mismatch gives the best oracle-discrepancy trajectory; these dynamics reflect a performance ceiling rather than downstream task performance.Source: §A.3.1 · pp. 18–19, Figure 10 · p. 20

A.3.2. ITERATION LABEL ANALYSIS

  • Continue-labeled tokens have over 5× the entropy of stop-labeled tokens, while tokens selected by both 1.7B and 4B references have about 2.0× the loss of non-overlap selected tokens.The 1.7B reference identifies 81% of the iteration-selected tokens found by the 4B reference, indicating a stable core of difficult token positions.Source: §A.3.2 · pp. 19–20, Figure 11 · p. 21

A.3.3. ITERATION DECIDER ANALYSIS

  • Across Math, Code, and QA, the decider yields 5.2–9.0% absolute accuracy gains over single-pass inference while being invoked on only 7.8–26.6% of tokens.At a fixed threshold, it iterates more often on QA (26.6%) than Math (7.8%), without retraining across domains.Source: §A.3.3 · p. 20

A.3.4. TOKEN ALTERNATION PATTERN

  • Figure 12 compares token cross-entropy distributions for tokens selected by both reference models and tokens selected by only one, across reference-model pairs.The caption describes the comparison but does not state which group has higher cross-entropy.Source: Figure 12 · p. 22

A.4.1. DUO-CAUSAL ATTENTION IMPLEMENTATION

  • Duo-causal attention lets queries attend to keys that are no later in token position or iteration depth, while preserving positional and iteration causality.It concatenates K/V across accessible depths and remains compatible with optimized attention kernels such as FlashAttention without kernel modifications.Source: §A.4.1 · pp. 23–24

A.3.5. DUO-CAUSAL ATTENTION PATTERN

  • TaH’s duo-causal attention shows layer-wise heterogeneity: lower layers allocate relatively less attention to first-iteration keys, while deeper layers show higher and varied cross-depth reuse.Representative heads attend mainly to first-iteration keys, mainly to second-iteration keys, or to both, demonstrating diverse learned cross-depth strategies.Source: §A.3.5 · pp. 22–23, Figure 15 · p. 24, Figure 14 · p. 24

A.5. Additional Related Work

  • Unlike layer-skipping methods that still load the full model’s parameters, TaH adds computational depth by iterating core tokens without increasing parameter count.The related-work discussion describes router-based and fixed layer-skipping approaches as dynamically bypassing layers.Source: §A.5 · p. 25

A.6. Limitations and Future Work

  • The comparison with official Qwen3 models is qualified by differences in training data, scale, and procedures; TaH uses SFT on limited public data.Future work proposes combining TaH with online distillation or reinforcement learning to improve supervision and optimize compute–accuracy trade-offs and continuation policies.Source: §A.6 · p. 25

A.7. Use of Large Language Models (LLMs)

  • The authors state that large language models were used only for proofreading and minor copyediting, not for scientific content, ideation, or method design.They attribute the research ideas, methods, and conclusions solely to the authors.Source: §A.7 · p. 25

A.8. Examples

  • In an AIME25 example, only TaH successfully answers the geometry question, while blue-highlighted tokens mark positions that undergo two iterations.The passage presents the example as an intuitive illustration of the method.Source: §A.8 · p. 25
Loading 2511.08577v4…