Can Small Models Overthink? TaH Skips 93% Redundant Iterations and Boosts Accuracy

Problem Motivation

Small, parameter‑constrained language models (e.g., 0.6 B–4 B) are attractive for edge deployment but often fail on reasoning tasks because a few critical tokens are predicted incorrectly, steering the entire reasoning chain off‑track. Existing Looped Transformers apply the same number of latent‑space iterations to every token, assuming that more computation always helps.

Latent Overthinking

Empirical analysis of Looped Transformers reveals a phenomenon called latent overthinking : many tokens are already correct after the first forward pass, yet additional latent iterations can flip them to wrong answers. Uniform extra iterations therefore produce both “fixes” and “breaks”.

Oracle Baseline

An oracle that re‑iterates only tokens that were wrong in the first pass improves downstream performance by up to 7.3 % while re‑iterating only 11 %–19 % of tokens. This establishes an upper bound for selective iteration.

Think‑at‑Hard (TaH) Architecture

TaH introduces three complementary components that enable token‑level dynamic latent iteration.

Lightweight iteration decider : a small MLP that consumes hidden states from the backbone after each latent iteration and outputs a continuation probability. If the probability falls below a preset threshold, decoding proceeds to the next token; otherwise another iteration is performed.

Duo‑causal attention : extends causal attention to a two‑dimensional grid of token position × iteration depth . For token i at depth d , the query can attend only to keys/values from earlier positions with depth ≤ d . This preserves full parallelism while allowing cross‑depth information flow.

Depth‑aware LoRA : LoRA adapters are activated only for iterations with depth > 1, focusing low‑rank adaptation on correcting difficult tokens while leaving the first‑pass next‑token prediction untouched. Cross‑iteration residual connections ensure deeper iterations refine the previous prediction instead of starting from scratch.

During inference TaH averages 1.07 iterations per token , effectively skipping ~93 % of second‑pass computation.

Two‑Stage Training Procedure

Because the decider’s decisions depend on the backbone’s prediction quality, end‑to‑end training is unstable. TaH adopts a decoupled two‑stage scheme:

Train the backbone with a static oracle that iterates only on tokens that were initially wrong.

Freeze the backbone and train the decider to mimic the oracle’s continue/stop decisions.

This approach dramatically improves convergence speed and training stability.

Experimental Setup

Backbones: Qwen3‑0.6B, Qwen3‑1.7B, Qwen3‑4B (pre‑trained on the Open‑R1 dataset). Evaluation benchmarks (nine total): GSM8K, MATH500, AMC23, AIME25, OlympiadBench, GPQA‑Diamond, MMLU‑STEM, HumanEval++, MBPP++.

Results

Without adding parameters, TaH improves accuracy by 3.0 %–3.8 % over the baseline.

TaH+ (adds ≤ 3 % parameters for decider, duo‑causal modules, etc.) raises gains to 5.3 %–6.2 % .

Compared with the prior Looped Transformer variant Ouro, TaH achieves 3.8 %–4.4 % higher scores; TaH+ reaches 6.1 %–6.8 % higher.

Average FLOPs increase by only 4 %–5 % . Memory usage drops by a factor of 1.48× and decoding speed improves by 2.48× relative to an always‑iterate baseline.

The oracle’s 7.3 % improvement (with 11 %–19 % token re‑iteration) is matched and exceeded when TaH’s architecture is applied.

Ablation Studies

Key findings from systematic ablations:

Replacing the dynamic decider with static depths (Always‑1 or Always‑2) degrades average performance by 6.1 % and 16.4 %, confirming the benefit of selective iteration.

Substituting duo‑causal attention with standard causal attention reduces scores by 5.4 %–8.5 %, highlighting the importance of cross‑depth attention.

Removing depth‑aware LoRA and cross‑iteration residuals lowers accuracy by 4.9 %.

Training the decider jointly with the backbone (single‑stage) leads to instability or collapse, validating the two‑stage training design.

Semantic Analysis of Decider Behavior

On the validation set, the tokens “But” and “So” trigger additional iterations most frequently, with probabilities of 34 % and 18 % respectively. These discourse markers often signal reasoning pivots, indicating that TaH learns to allocate extra compute to semantically important positions.

Resources

Code repository: https://github.com/thu-nics/TaH

Project page: https://fuvty.github.io/TaH_project_page

Paper (PDF): https://arxiv.org/pdf/2511.08577

Conclusion

TaH demonstrates that fine‑grained, token‑level dynamic compute allocation can outperform brute‑force scaling for small reasoning models. By combining a lightweight decider, duo‑causal attention, and depth‑aware LoRA within a stable two‑stage training pipeline, TaH achieves consistent accuracy gains while reducing compute and memory overhead. The work opens a path toward test‑time scaling techniques that adapt compute based on token difficulty.

Code example

[1] Jaech, A., Kalai, A., Lerer, A., et al. OpenAI o1 system card. arXiv preprint arXiv:2412.16720, 2024.
[2] Guo, D., Yang, D., Zhang, H., et al. DeepSeek-R1: Incentivizing reasoning capability in LLMs via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025.
[3] Yang, A., Li, A., Yang, B., et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025.