How Should Reasoning Be Organized in a Transformer's Latent Space?

Hongyu Gu; Chang Liu; Jingwen Fu

cs.AISep 12, 2026

How Should Reasoning Be Organized in a Transformer's Latent Space?

Authors: Hongyu Gu, Chang Liu, Jingwen Fu

Organizations: University of Science and Technology of China Hefei, China · Zhongguancun Academy Beijing, China

Abstract

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.

Figures & tables

Figure 1: The cumulative-state hypothesis from latent state to computation. (a) The frontier state keeps only the states active now, whereas the cumulative state keeps all reached states for the same final task. (b) Each update adds new states; equalizing the weights keeps every reached state available. (c) Equal cumulative weights avoid a weakly represented state and are optimal when future use is unknown.

Figure 2: Why remembering more can require less width. (a) Across many ways of using the reached history state, cumulative state stays below the predicted width bound, while the deliberately hardest examples approach it. (b) As the reached history grows, cumulative keeps useful signal for the final task, whereas frontier discards earlier states and becomes harder to use.

Figure 3: Trained Transformers reproduce the predicted width advantage. The shared legend shows six ways to weight the same reached history. (a1) Results for the two-layer block model; (a2) results for the GPT-2-block model, with bands showing variation across runs. Every curve in a panel uses the same trained attention module, and equal cumulative weighting reaches

90\%

accuracy on the hardest tested task first. (b) Five paired runs confirm that cumulative reaches the same held-out accuracy at a smaller width than frontier in both architectures. Here

d_{0.9}

means the smallest interpolated width reaching

90\%

accuracy; Ideal supplies the intended state directly.

Figure 4: Why equal cumulative weights are safest. (a) Each point compares the width needed to keep states usable during reasoning and the width needed to use them for the final answer; lower is better. (b) The heatmap shows how much state is assigned to each position in the reached history. (c) Using the same trained attention module, states with small weights are the first to be missed. The summary reports success on the whole routing step and the rate at which inactive states are selected by mistake.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Split	# traces	$\|V\|$	$\|E\|$	solution length
Train	14,785	22.8	36.5	3.5
Validation	257	22.7	36.3	3.5
Test	419	22.7	36.0	3.5

Appendix

Table 1: ProsQA graph-4 statistics. Graph sizes and solution lengths are averaged over traces.

Split	Used traces	Queries	Support fraction	Realized $C_{\max}$
Train	14,785	118,280	$[0.750,0.857]$	1.319
Validation	120	480	$[0.750,0.857]$	1.271
Test	419	3,352	$[0.750,0.857]$	1.287

Appendix

Table 2: Audit of the fixed nonlocalized query bank. “Queries” counts coefficient-space queries before repeated Gaussian test draws.

Component	Configuration
Widths	$16,24,32,40,48,56,64,80,96,112,128$
Training seeds / embedding seeds	$0{:}4$ / $5100{:}5104$
Query-bank size per trace (train/val/test)	$8/4/8$
Decoys (train/val/test)	$256/1{,}024/4{,}096$
Gaussian draws (val/test)	$1/4$
Epochs and optimizer	4; AdamW, learning rate $3\times 10^{-3}$

Appendix

Table 3: Configuration of the lightweight two-block paired experiment.

$d$	Learned frontier	Learned cumulative	Ideal frontier	Ideal cumulative
16	$0.048\pm 0.001$	$0.254\pm 0.012$	$0.052\pm 0.001$	$0.442\pm 0.003$
24	$0.114\pm 0.004$	$0.574\pm 0.008$	$0.120\pm 0.005$	$0.790\pm 0.004$
32	$0.210\pm 0.007$	$0.793\pm 0.015$	$0.216\pm 0.007$	$0.947\pm 0.005$
40	$0.322\pm 0.005$	$0.899\pm 0.010$	$0.328\pm 0.006$	$0.990\pm 0.002$
48	$0.430\pm 0.008$	$0.946\pm 0.004$	$0.436\pm 0.008$	$0.9987\pm 0.0002$
56	$0.533\pm 0.010$	$0.972\pm 0.005$	$0.540\pm 0.009$	$0.9999\pm 0.0002$

Appendix

Table 4: Held-out top-1 accuracy at every tested width.

Seed	Frontier	Cumulative	C/F
0	108.9	39.2	0.360
1	112.1	39.9	0.356
2	110.3	40.4	0.366
3	111.7	41.6	0.372
4	110.1	40.6	0.369
Mean	110.6	40.3	0.365

Appendix

Table 5: Interpolated empirical critical dimension

d_{0.9}

for every lightweight-model training seed.

Component	Configuration
Backbone	causal GPT-2; 2 layers, 8 heads, FFN width $4d$
Transformer details	GELU, context 512, dropout 0.1, residuals and LayerNorm
Tested widths / maximum codebook width	$32,64,96,128,192,256,384$ / 768
Model seeds / embedding seeds	$0{:}4$ / $6100{:}6104$
Used traces (train/validation/test)	$800/120/120$
Queries per trace (train/validation/test)	$8/4/4$

Appendix

Table 6: Complete configuration of the paired GPT-2 width sweep.

$d$	Frontier	Cumulative
32	$0.468\pm 0.038$	$0.898\pm 0.024$
64	$0.900\pm 0.026$	$0.998\pm 0.003$
96	$0.978\pm 0.009$	$1.000\pm 0.000$
128	$0.993\pm 0.004$	$1.000\pm 0.000$
192	$0.998\pm 0.002$	$1.000\pm 0.000$
256	$1.000\pm 0.001$	$1.000\pm 0.000$

Appendix

Table 7: Complete GPT-2 held-out top-1 accuracy sweep.

Seed	Frontier	Cumulative	C/F
0	63.2	32.0	0.507
1	68.6	33.0	0.481
2	73.5	36.2	0.493
3	62.6	32.0	0.511
4	63.0	38.8	0.616
Mean	66.2	34.4	0.521

Appendix

Table 8: GPT-2 interpolated critical dimension

d_{0.9}

for every model seed.

Codebook seed	Uniform cumulative	Recency-heavy	Random mild
4903	$0.896/0.032$	$0.576/0.572$	$0.533/1.318$
5004	$0.904/0.028$	$0.594/0.585$	$0.541/1.282$
5105	$0.907/0.028$	$0.593/0.557$	$0.536/1.312$
5206	$0.913/0.026$	$0.587/0.564$	$0.541/1.336$
5307	$0.909/0.029$	$0.578/0.570$	$0.534/1.280$
Mean	$0.906/0.029$	$0.586/0.570$	$0.537/1.306$

Appendix

Table 9: Raw reusable-token routing results at

d=192

. Each cell reports mean joint route-step success (fraction) / per-inactive-key false-positive rate (%). A route step succeeds iff every currently reached node is accepted and every unreached graph node and all 512 Gaussian decoys are rejected.

Explore similar work

State Stream Transformer (SST) V2: Parallel Training of Nonlinear Recurrence for Latent Space Reasoning
Apr 30, 2026Thea AvissEfficient Latent Reasoning Recurrent Neural Networks
Fifth Dimension
Reasoning Primitives in Hybrid and Non-Hybrid LLMs: Do Architectural Differences Yield Advantages in State-Tracking and Recall?
Apr 23, 2026Shivam Rawat, Lucie Flek, Florian Mai +1LLM Reasoning Strategies Search-Augmented Reasoning
Lamarr Institute for Machine Learning and Artificial Intelligence · Rheinische Friedrich-Wilhelms-Universität Bonn
Latent Thought Flow: Efficient Latent Reasoning in Large Language Models
Jun 15, 2026Xiandong Zou, Jing Huang, Jianshu Li +1Efficient Latent Reasoning Latent Thoughts
Singapore Management University · Ant Group

Apr 30, 2026, cs.LGJ/K move · Enter open · S save

State Stream Transformer (SST) V2: Parallel Training of Nonlinear Recurrence for Latent Space Reasoning

Thea Aviss

Fifth Dimension

Efficient Latent Reasoning Recurrent Neural Networks Deliberation

Open paper AskPDF arXiv

Current transformers discard their rich latent residual stream between positions, reconstructing latent reasoning context at each new position and leaving potential reasoning capacity untapped. The State Stream Transformer (SST) V2 enables parameter-efficient reasoning in continuous latent space through an FFN-driven nonlinear recurrence at each decoder layer, where latent states are streamed horizontally across the full sequence via a learned blend. This same mechanism supports continuous latent deliberation per position at inference time, dedicating additional FLOPs to exploring abstract reasoning before committing to a token. A two-pass parallel training procedure resolves the sequential dependency of the recurrence to allow compute-efficient training. Hidden state analysis shows the state stream facilitates reasoning through exploration of distinct semantic basins in continuous latent space, where transitions at content-dependent positions move the model into a substantially different Bayesian posterior, directly influencing the latent space at future positions. We also find, via a learned probe, that at the first generated token position, the latent state already predicts whether the eventual answer will survive or break under additional latent computation for every subsequent position. Co-trained into an existing 27B backbone using only a small dataset of GSM8K examples, the SST delivers a +15.15 point gain over a fine-tuning-matched baseline on out-of-distribution GPQA-Diamond and cuts that same baseline's remaining GSM8K errors by 46%, together showing that the reasoning improvement is attributable to the architectural mechanism rather than scale or training data. On GPQA-Diamond, the resulting 27B SST also achieves higher accuracy than several larger open-weight and proprietary systems, including open-weight models up to 25 times larger.

Abstract

Explore similar work

State Stream Transformer (SST) V2: Parallel Training of Nonlinear Recurrence for Latent Space Reasoning

Reasoning Primitives in Hybrid and Non-Hybrid LLMs: Do Architectural Differences Yield Advantages in State-Tracking and Recall?

Latent Thought Flow: Efficient Latent Reasoning in Large Language Models