Think-at-Hard — Selective Latent Iteration for Looped Reasoning Transformers

Paper. Tianyu Fu, Yichen You, Zekai Chen, Guohao Dai, Huazhong Yang, Yu Wang. Think-at-Hard: Selective Latent Iterations to Improve Reasoning Language Models. ICML 2026. Tsinghua University (NICS-EFC Lab). [arXiv]

0. The Picture in One Paragraph

Looped transformers give every token a second (or third) latent forward pass before it commits to an output — a cheap, training-data-free form of test-time reasoning. The implicit assumption is that more iteration is always at least neutral. Think-at-Hard (TaH) shows this assumption is false: a fixed “always iterate twice” policy flips more correct first-pass predictions into wrong ones than it fixes — a latent overthinking phenomenon, the looped-transformer analogue of overthinking in explicit chain-of-thought. TaH replaces the fixed policy with a lightweight neural decider that predicts, per token, whether the first pass is likely wrong and only then triggers a second latent iteration. To make this selective computation actually fast in parallel hardware, TaH introduces duo-causal attention — a causal mask extended across a 2D (sequence position, iteration depth) grid — plus depth-aware LoRA adapters that activate only on the extra pass. The result: TaH beats a fixed-two-iteration baseline by 8–11% while skipping 94% of second iterations, and beats a same-data single-iteration Qwen3 baseline by roughly 4–5%.

1. The Problem — Latent Overthinking

A looped transformer processes a token through a shared block once, then optionally again, refining the hidden state before it is decoded. The implicit bet made by every prior looped-transformer design (Universal Transformers, Huginn-style recurrent depth) is that iterating uniformly across all tokens can only help or be neutral.

TaH’s diagnostic experiment — an “AlwaysThink” ablation that forces a second iteration on every token — shows the opposite: tracking per-token prediction transitions across iteration depths reveals that more predictions flip from correct to incorrect than from incorrect to correct. Most tokens are already right after the first pass; the second pass, applied indiscriminately, does more harm than good on average. This mirrors the well-documented overthinking failure mode in explicit CoT reasoning, but at the level of a single token’s hidden state rather than a whole reasoning trace.

2. The TaH Framework

2.1 Lightweight neural decider

A small classifier reads the first-pass hidden state at each token position and predicts whether that token’s first-pass top-1 prediction is likely wrong. It is trained against oracle labels: a token is labeled “iterate” if running the second pass actually corrects it relative to a reference. At inference, the decider gates the second iteration per token — most tokens skip it entirely.

2.2 Duo-causal attention

Selective, per-token iteration creates a structural problem: different tokens in the same sequence are now at different iteration depths, but training and inference still need to run in parallel, not token-by-token. TaH’s fix is to extend the causal mask from the usual 1D (sequence position) axis to a 2D grid of (position, depth): a token attends to all earlier positions and to states at or below its own iteration depth at the same and earlier positions. This preserves causality on both axes while keeping full sequence-level parallelism — no serialization penalty for mixing 1-pass and 2-pass tokens in one forward call.

2.3 Depth-aware LoRA adapters

Small LoRA adapters are active only during the second (extra) pass. This decouples the backbone’s original next-token objective (used on the first pass) from a narrower refinement objective (used only when the decider triggers a second pass), letting the extra parameters specialize in fixing likely-wrong tokens rather than diluting the general-purpose first-pass behavior.

2.4 TaH+ — a cheap upper bound

A variant, TaH+, skips the decider and its selective gating entirely, always running the second pass with the LoRA adapters active (under 3% extra parameters). It is not meant to be efficient — it is a cheap way to check how much of TaH’s decider-gated gain comes from the adapters themselves versus the selectivity.

3. Training

TaH is fine-tuned from Qwen3-Base at 0.6B and 1.7B scale on the Open-R1 reasoning dataset. The decider is trained as a binary classifier against oracle iterate/don’t-iterate labels derived from per-token correctness deltas between the first and second pass. Evaluation spans nine benchmarks across math, QA, and coding, including GSM8K, MATH500, and AIME24/AIME25.

4. Results

Comparison	TaH gain	Second-iteration compute saved
vs. AlwaysThink (fixed 2-pass on every token)	+8.1 to +11.3%	94% of second passes skipped
vs. Ouro (looped-transformer baseline, same data)	+3.8 to +4.4%	93% of second passes skipped
vs. Standard single-iteration Qwen3 (same data)	~4–5%	—
Oracle iteration policy (upper bound)	up to +7.3%	—

On AIME25, TaH reaches 17.9% versus a 13.3% standard single-iteration baseline. The gap between TaH’s realized gains and the oracle ceiling of +7.3% is itself informative: the trained decider captures most, but not all, of the available headroom — there is still knowable-but-unexploited signal about which tokens need a second pass.

5. Why It Matters

Looped/recurrent-depth transformers (this blog has covered several — see the recurrent-depth and elastic-depth family below) treat “how many iterations” as either a fixed hyperparameter or a global, sequence-level budget. TaH pushes the granularity down to individual tokens, and — more importantly — demonstrates empirically that uniform iteration is not merely wasteful but actively harmful on net. That reframes the design question for the whole looped-transformer line: the goal isn’t just “skip iterations to save compute,” it’s “skip iterations because most of them are actively hurting accuracy.”

6. Limitations Worth Knowing

Decider imperfectly approximates the oracle. The 7.3% oracle ceiling versus TaH’s realized 8–11% gain over AlwaysThink (a different, weaker baseline) shows there is still room for a better decider — the selectivity mechanism is good but not optimal.
Small-scale evaluation. Results are reported at 0.6B and 1.7B parameter scale; whether the overthinking phenomenon and the decider’s effectiveness hold at larger scales is untested here.
Binary, per-token gating is coarse. The decider makes a single iterate/don’t-iterate call; it does not decide how much extra computation a hard token deserves beyond one more pass.
Fine-tuned rather than pretrained from scratch. TaH is built on top of an existing Qwen3-Base checkpoint; interactions between the decider and pretraining-time objectives are not explored.

7. The Takeaway for a First Reader

If you remember three things:

Looped transformers can overthink. Forcing a second latent iteration on every token flips more correct predictions to incorrect than the reverse — indiscriminate iteration is a net negative, not just a compute cost.
TaH fixes this with three pieces: a lightweight decider that gates the second pass per token, duo-causal attention that lets mixed-depth tokens be processed in parallel, and depth-aware LoRA adapters that specialize the extra pass toward refinement.
8–11% accuracy gain over a fixed-iteration baseline while skipping 94% of second passes — selective computation beats uniform computation, and the decider captures most (not all) of the oracle’s theoretical headroom.

References

Fu, T., You, Y., Chen, Z., Dai, G., Yang, H., & Wang, Y. (2026). Think-at-Hard: Selective Latent Iterations to Improve Reasoning Language Models. ICML 2026. arXiv:2511.08577.
Related on this site: LoopFormer — elastic, budget-conditioned loop depth trained via shortcut-consistency, a complementary approach to token-level selective iteration; BG-MCTS — another budget-aware test-time mechanism, at the level of search rather than latent iteration.