Source-linked AI summary
How Should Reasoning Be Organized in a Transformer's Latent Space?
Hongyu Gu, Chang Liu, Jingwen Fu
TL;DR
The paper asks what a latent state should preserve as continuous reasoning unfolds, since prior work shows that superposition can represent parallel alternatives but leaves this design choice open. It compares frontier and cumulative state weightings and finds that equal-weight cumulative states can support attention at smaller width, with experiments supporting the predicted advantage and failure pattern.
Problem
Prior work establishes that a latent state can represent parallel search alternatives, but leaves open what it should preserve as computation continues.
Source: §1 · p. 1Method
The paper compares frontier, equal-weight cumulative, and general weighted states in a recurrent Transformer setting and analyzes their attention and width requirements.
Source: §3 · pp. 3–4Results
Cumulative states reach the same accuracy at smaller width than frontier states in lightweight and GPT-2-block Transformers, while unequal weights fail first on least-weighted states.
Source: §7 · p. 9Takeaways & Limitations
When past reasoning may be useful again, preserve reached states with equal weights and restore that balance as computation proceeds.
Source: §7 · p. 9The comparison assumes orthonormal competing keys orthogonal to the useful direction and state span, with a common bank restricted by M≤n−m.
Source: §3 · p. 4
Abstract
from arXiv · showhide
Continuous reasoning has emerged as a promising way to improve reasoning in large language models (LLMs). Yet we still lack a clear principle for deciding what a latent state should preserve. Reasoning by superposition shows that a single latent state can encode several search alternatives and expand them in parallel. We ask how those states should be weighted as reasoning proceeds. A natural choice is to preserve only the states active at the frontier step, since keeping every reached state appears to spread a limited hidden width too thin. We show that the opposite can hold. When later computation draws on several reached states, a cumulative state can guide attention correctly at a smaller hidden width than a frontier state that stores fewer states. At the same width, the cumulative state therefore keeps more intermediate states available for later reasoning. More generally, equal cumulative weights are optimal when future queries are unknown and remain close to the best task-specific weights when those queries are known. Experiments with two-layer and GPT-2 Transformers reproduce the predicted width advantage and show that unequal weights fail first on the states that receive the least weight. This suggests a important principle: keep reached states equally weighted, and restore equal weights as computation proceeds.
1 INTRODUCTION
The section asks what a latent state should preserve as reasoning unfolds, beyond the established ability of superposition to represent parallel alternatives. It argues that retaining all reached states equally can support later attention at smaller width, and presents a rule for unknown or known future queries.
1 INTRODUCTION
- The comparison evaluates frontier, equal-weight cumulative, and general nonnegative weighted states using the same task, candidate answers, Transformer, and final query.Source: §1 · p. 1
- Cumulative states can guide later attention at smaller hidden width than frontier states because useful components reinforce while unrelated candidates remain random interference.The comparison holds when later computation draws on several reached states.Source: §1 · pp. 1–2
- Figure 1 contrasts frontier and cumulative retention, shows equal-weight updates, and depicts why equal cumulative weights avoid weakly represented states when future use is unknown.Source: Figure 1 · p. 2
- Equal cumulative weights are optimal when future use is unknown and require at most a C2 factor of the best task-specific width when queries are known.The unknown-query guarantee applies when only the stated coefficient bound is known.Source: §1 · pp. 1–2
- The Transformer analysis tracks attention retrieval, value writes, residual retention, and output-projection weighting across all D steps, identifying the least represented state as the bottleneck.Source: §1 · p. 2
- Two-layer and GPT-2 experiments support the width advantage, and unequal weights fail first on states assigned the least weight.Tests include synthetic query families, paired width sweeps, and position-level interventions.Source: §1 · p. 2
2 RELATED WORK
Prior work establishes parallel search through frontier superposition and studies feature capacity, recurrent computation, and latent-state reasoning. This paper connects those topics to later attention over reached states and the Transformer operations that update them.
2 RELATED WORK
- Zhu et al. (2025b) show that frontier superposition keeps current alternatives active and expands them together for parallel search; this work studies capacity across a complete D-step computation.Source: §2 · pp. 2–3
- Feature-superposition analyses study feature representation and recoverability, while this work asks whether normalized latent states support later attention over multiple reached states.When later attention combines reached states, aligned components add coherently while unrelated candidates remain random interference.Source: §2 · p. 3
- Recurrent and buffer-based models carry information through latent states, and this analysis links attention retrieval, value writes, and residual retention to the least represented state's bottleneck.It yields equal weighting of reached states when future use is unknown.Source: §2 · p. 3
3 PROBLEM SETTING AND CENTRAL HYPOTHESIS
The paper models a recurrent two-block Transformer that retrieves objects with attention and writes their values into a continuous thought. It compares frontier retention, equal-weight cumulative retention, and general nonnegative history weights under bounded future-query assumptions.
3 PROBLEM SETTING AND CENTRAL HYPOTHESIS
- A recurrent two-block Transformer uses QK⊤ attention to choose active objects, then writes their values into the next continuous thought.Object embeddings are available to attention as keys and values through the copied-edge interface.Source: §3 · p. 3
- At depth c, F contains states first reached then and R contains states reached up to then, with k=|F| and m=|R|.Vertices may represent intermediate facts, entities, subgoals, partial plans, or candidate states.Source: §3 · p. 3
- The proposed state normalizes the embedding mixture for a coefficient vector, while frontier, uniform cumulative, and general weighted states respectively retain current states, all reached states equally, or all reached states with nonnegative weights.Source: §3 · pp. 3–4
- The final attention head maps the thought to a query and candidate tokens to keys and values; success means the useful value receives the largest attention weight.The competing keys are assumed orthonormal and orthogonal to the useful direction and state span, with a common bank requiring M≤n−m.Source: §3 · p. 4
- Future useful directions are nonnegative, unit-normalized, supported on reached states, and bounded so no single state receives more than C/√m weight.C measures how unevenly a future computation may use the reached history; uniform queries have C=1.Source: §3 · p. 4
- The hypothesis predicts that uniform cumulative weights are within C2 of the best sequence chosen with advance query-family knowledge and are the best single choice when only the coefficient bound is known.Source: §3 · p. 4
4 THE WIDTH OF A D-STEP TRANSFORMER COMPUTATION
This section defines width by whether a weighted state sequence supports the required attention calls and final answer across a D-step computation. It relates that width measure to repeated Transformer updates, with a stability guarantee under stated conditions.
4 THE WIDTH OF A D-STEP TRANSFORMER COMPUTATION
- The state sequence must keep designated sources usable at each step and meet the correct final-answer event with probability at least 1−δ for every query in U.The width is defined for a coefficient sequence whose weights at each step belong to the allowed history weights.Source: §4 · pp. 4–5
- For fixed D, δ, and positive final-query alignment, the required joint width for D-step computation scales as Θ(D(a; U)).The guarantee assumes Θ(n) inactive node keys at each step and a common final bank.Source: §4 · p. 5
- For cumulative or frontier updates, sufficiently large width makes the actual D-step computation and final answer correct with probability at least 1−δ for every query in U.Finite attention scales suffice, with the width constant depending on fixed D and δ.Source: §4 · p. 5
- Repeated writing and normalization errors can remain below attention gaps across the trace, while the smallest retained coefficient controls next-step routing and final-query alignment.The destination-write error can be made arbitrarily small by finite attention scales.Source: §4 · p. 5
5 THEORETICAL VALIDATION OF CUMULATIVE STATE
The analysis compares uniform cumulative weights with task-specific weightings and frontier states. It gives a C2 guarantee for query classes, identifies when preserving history reduces width, and describes a Transformer construction for both state choices.
5.1 THE C2 GUARANTEE FOR A KNOWN QUERY CLASS
- For a known query class, equal cumulative weights require at most C2 times the width of the best task-specific weights; when the query is unknown, cumulative is best in the worst case.The comparison allows every nonnegative coefficient sequence, including sequences a scalar residual update cannot generate.Source: §5.1 · p. 6
5.2 WHY PRESERVING HISTORY CAN REDUCE WIDTH
- For a fixed C, the growing-history regime k = o(m) includes chains, fixed-width traces, and grid-like breadth-first search.The stated query condition also covers whole-set aggregation, additive objectives over visited objects, and supervision over algorithmic traces.Source: §5.2 · p. 6
- For history-wide computation, cumulative uses a factor m/k less width than frontier; Theorem 4.1 combines this gain with the cost of producing preceding states.The end-to-end gain holds when final answer selection is the active bottleneck for both policies; for fixed C, the ratio vanishes as k = o(m).Source: §5.2 · p. 6
- Cumulative history helps because useful historical components add before normalization while unrelated candidates remain centered fluctuations.The query condition covers whole-set aggregation, additive route costs, algorithms, and supervision over algorithmic traces.Source: §5.2 · pp. 6–7
5.3 HOW THE SAME TRANSFORMER WRITES FRONTIER OR CUMULATIVE
- The Transformer construction uses gates (0, 1) for frontier and (√mc, √rc) for cumulative, approximating uniform cumulative states.The update combines a scaled copy of the prior state with destination values through a fixed output projection and RMSNorm.Source: §5.3 · p. 7
6 EXPERIMENTAL EVIDENCE
Across synthetic queries and trained Transformers, cumulative states reach target accuracy at smaller widths than frontier states. Unequal weighting leaves low-weight states most vulnerable to missed routing.
6 EXPERIMENTAL EVIDENCE
- Within each experiment, the task, candidates, embedding codebook, and trained attention mechanism are shared; only the imposed state coefficients change.The protocol compares state constructions under controlled conditions.Source: §6 · p. 7
6.1 TEST I: DOES CUMULATIVE NEED LESS WIDTH?
- Across synthetic query families, cumulative state stays below the predicted width bound, while deliberately hardest examples approach it as the reached history grows.The tests include dense, temporal, block, union, sharp-witness, and zero-overlap query constructions.Source: §6.1 · p. 7, Figure 2 · p. 8
6.2 TEST II: DO TRAINED TRANSFORMERS SHOW THE SAME ADVANTAGE?
- 40.3 ± 1.1 versus 110.6 ± 1.6 in the lightweight two-block model, and 34.4 ± 3.8 versus 66.2 ± 5.9 with the two-layer GPT-2 backbone, are the cumulative and frontier widths for 90% held-out success.Both comparisons come from paired five-seed sweeps; Figure 3 reports the width sweeps and paired minimum-width estimates.Source: Figure 3 · p. 8, §6.2 · p. 8
6.3 TEST III: WHAT FAILS UNDER UNEQUAL WEIGHTS?
- Joint route success falls from 0.906 with uniform weights to 0.586 with recency weights and 0.537 with random weights; per-inactive-key false-positive rates rise from 0.03% to 0.57% and 1.31%, respectively.Low-coefficient tokens receive lower membership margins and are the first to be missed; Figure 4 depicts the weighting and routing comparisons.Source: Figure 4 · p. 9, §6.3 · p. 9
7 CONCLUSION
The conclusion recommends preserving all reached states with equal weights and restoring that balance during computation. The paper reports that this principle is supported by theory and experiments in lightweight and GPT-2-block Transformers.
7 CONCLUSION
- When past reasoning may be useful again, preserve reached states with equal weights rather than keeping only the latest frontier.The conclusion attributes the width advantage to useful history components reinforcing one another in later attention, with the residual retaining earlier information and the feed-forward layer restoring equal weights after each write.Source: §7 · p. 9
AI USE STATEMENT
- The authors report using generative AI tools for grammar correction, writing optimization, and some chart-drawing code, with AI-assisted content reviewed by the author.The author states responsibility for the final content, including proofs, experiments, citations, charts, and text.Source: AI USE STATEMENT · p. 10
A PROOFS FOR THE DIMENSION LIMIT
This appendix develops local attention-comparison facts used in the complete-trace theorem. It gives a distributional representation, a monotonicity result, and proof steps for a local transition.
A PROOFS FOR THE DIMENSION LIMIT
- For a unit query direction and aligned state with alignment μ, the attention comparison is represented using independent Gaussian variables and a χ-distributed variable.The representation is the stated exact distribution for one attention comparison.Source: A PROOFS FOR THE DIMENSION LIMIT · p. 12
- For fixed dimension and decoy count, success probability increases strictly with alignment μ, yielding the critical-dimension conclusion.The proof orders the aligned score stochastically and notes that conditional success against M decoys, Φ(s)^M, increases with that score.Source: A PROOFS FOR THE DIMENSION LIMIT · p. 12
- The local transition proof uses Gaussian-maximum and χ concentration, then brackets the true score and maximum decoy on opposite sides of λ when λ is above or below 1.The proof concludes the transition statement (41) from these comparisons.Source: A PROOFS FOR THE DIMENSION LIMIT · p. 13
- M_max is defined as the largest candidate count M≤N meeting the target success probability; the appendix states that solving the transition relation for log M gives (58).The stated condition for the additional consequence is that min{dμ², log N} tends to infinity.Source: A PROOFS FOR THE DIMENSION LIMIT · p. 14
B PROOFS FOR FRONTIER, CUMULATIVE, AND WEIGHTED STATES
The proofs establish width guarantees for recurrent frontier and cumulative states, then show that cumulative weighting is robustly optimal and its C² comparison factor can be sharp.
B PROOFS FOR FRONTIER, CUMULATIVE, AND WEIGHTED STATES
- For a fixed query class, worst-query success is governed by its minimum alignment Γ, which determines the critical dimension.The proof applies the Gaussian critical-dimension law to the worst-query alignment and combines routing and final-answer requirements.Source: B PROOFS FOR FRONTIER, CUMULATIVE, AND… · pp. 14–15
- Successful routing requires the weakest source at each step to beat Θ(n) inactive keys, giving a matching lower-width bound.The lower bound scales as Ω(log n/a_min,c²) for the weakest source alignment.Source: B PROOFS FOR FRONTIER, CUMULATIVE, AND… · p. 15
- Controlling state error preserves routing and final-attention margins in the actual Transformer, yielding the sufficient-width guarantee for fixed depth.The proof bounds error accumulation across steps and transfers ideal attention gaps to the recurrent computation.Source: B PROOFS FOR FRONTIER, CUMULATIVE, AND… · p. 16
- The C² factor is sharp: a task-specific terminal state can attain alignment one while cumulative weighting has alignment 1/C.The construction attains equality in the width comparison under the stated readout-scale condition.Source: B PROOFS FOR FRONTIER, CUMULATIVE, AND… · p. 16
- Cumulative weighting minimizes the robust objective, and equal cumulative weights need at most a C² width factor versus the best task-specific weighting.The comparison combines bounds on routing and readout terms; the implementation guarantee separately accounts for construction stability.Source: B PROOFS FOR FRONTIER, CUMULATIVE, AND… · pp. 16–17
C EXACT REDUCTION FOR AN ARBITRARY FIXED BANK
This block extends the Gaussian reduction to arbitrary fixed candidate banks and documents the query construction and paired lightweight-model sweep.
C EXACT REDUCTION FOR AN ARBITRARY FIXED BANK
- For an arbitrary fixed bank, critical dimension follows the squared typical scale of a Gaussian maximum; the simplified log M/µ² form requires common competition geometry.The reduction needs neither semantic query selection nor orthogonality; correlation changes the competing Gaussian maximum.Source: C EXACT REDUCTION FOR AN ARBITRARY… · p. 17
D.1 DATASET AND QUERY CONSTRUCTION
- The ProsQA graph-4 queries include the complete terminal frontier and 75%–85.7% of reached history, with coefficient ratio at most 1.2.Frontier and cumulative reuse the same query, Gaussian draw, useful candidate, and decoy bank.Source: §D.1 · p. 18
D.2 LIGHTWEIGHT TWO-BLOCK WIDTH SWEEP √
- The lightweight two-block sweep shares router parameters and optimizer updates, tests widths 16–128, and evaluates top-1 retrieval against 4,096 decoys.The paired comparison uses five training seeds and held-out traces.Source: §D.2 · p. 18
D.3 COMPLETE FIXED-QUERY RESULTS
The fixed-query results show a lower critical width for cumulative states, even though their learned-state cosine is lower at width 40.
D.3 COMPLETE FIXED-QUERY RESULTS
- At width 40, frontier cosine is 0.966 ± 0.006 versus cumulative cosine 0.933 ± 0.006, so higher state cosine does not explain cumulative’s accuracy advantage.The threshold was stable at the individual-seed level.Source: §D.3 · p. 19, Table 4 · p. 19
- 0.365 ± 0.008: learned cumulative/frontier critical-dimension ratio, versus the fixed-alignment prediction of 0.318.The analytic cumulative and frontier states first crossed the tested grid at widths 32 and 112, respectively.Source: Table 5 · p. 19, §D.3 · p. 19
D.4 GPT-2 REPLICATION: CONFIGURATION AND COMPLETE RESULTS
The GPT-2 replication retains cumulative’s lower critical width across paired seeds; weighted-routing results also track the weakest coefficient, and the appendix reports computing resources.
D.5 WEIGHTED ROUTING PROTOCOL AND RAW CODEBOOK RESULTS
- The weakest coefficient’s normalized mean is 1.000, 0.814, and 0.775 across the compared conditions; it tracks joint success and false positives monotonically.This token-level pattern matches the predicted a_min mechanism.Source: §D.5 · p. 20, Table 9 · p. 21
D.4 GPT-2 REPLICATION: CONFIGURATION AND COMPLETE RESULTS
- 0.521 ± 0.067: mean cumulative/frontier critical-dimension ratio in the GPT-2 replication, with cumulative lower for every paired seed.The replication uses a two-layer causal GPT-2 Transformer and evaluates against 4,096 fixed decoys.Source: §D.4 · p. 19, Table 8 · p. 20
D.6 SOFTWARE AND COMPUTING RESOURCES
- The experiments used Python 3.10, PyTorch 2.5.1, CUDA 12.4, and NVIDIA A100 GPUs, with 55 lightweight models and 35 GPT-2 runs.The four-width weighted run used 0.17 GPU-hours.Source: §D.6 · p. 21