Papers
Topics
Authors
Recent
Search
2000 character limit reached

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

Published 30 Apr 2026 in cs.LG and cs.CL | (2605.00206v1)

Abstract: 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.

Authors (1)

Summary

  • The paper introduces an innovative two-pass training method that parallelizes nonlinear recurrence to enable efficient propagation of latent states.
  • The model leverages per-layer learned blend coefficients to orchestrate content-sensitive state updates, revealing emergent basin transitions and adaptive halting signals.
  • Robust numerical results on GPQA-Diamond and GSM8K benchmarks demonstrate significant accuracy gains with a 27B model, underscoring its parameter efficiency.

State Stream Transformer V2: Parallel Nonlinear Recurrence for Latent Space Reasoning

Architectural Overview

State Stream Transformer (SST) V2 introduces a latent space reasoning architecture that augments standard transformers with a per-layer, non-linear horizontal recurrence—the state stream. Each decoder layer maintains a persistently evolving latent state, updated via its own feedforward network (FFN) and propagated horizontally across the token sequence using learned, per-dimension blend coefficients. Unlike previous architectures (e.g., Universal Transformers, Mamba, Coconut), SST V2 combines vertical computation depth (iterative refinement at each position) with horizontal temporal continuity during token generation, enabling parameter-efficient computation by leveraging the pretrained FFN weights instead of increasing model capacity. The blending mechanism is non-bypassable, and per-layer blend strengths are learned independently, ensuring content-sensitive propagation.

Parallel Training of Sequential Recurrence

A major challenge with nonlinear, cross-position recurrence is parallelization. SST V2 resolves this using a two-pass training method: the first pass disables the blend, producing standalone FFN outputs for each layer and position; these are then propagated across positions via an associative scan, yielding approximate states. The second pass re-enables the blend with these scan states, allowing parallel processing across all sequence positions. The approximation error is O(α2)O(\alpha^2) in the blend coefficient, and co-training of adapters, blend parameters, and normalization weights ensures the model learns to maximize utility across both passes. This method enables practical fine-tuning on large pretrained transformers, making nonlinear latent recurrence feasible without prohibitive computational cost.

Mechanism and Dynamics of Latent Computation

SST V2 fundamentally changes the internal computation of the model, leading to emergent latent dynamics in the hidden state space. Analysis reveals "punctuated equilibrium": most positions exhibit stable hidden representations across iteration depths, but specific, content-dependent positions trigger sharp reorganizations spanning all layers. Top-kk overlap metrics confirm a bimodal regime—stable (mean overlap $0.990$) versus low-overlap (mean $0.869$)—where the latter induces dramatic shifts in the Bayesian posterior. These basin shifts affect future computation, causing large argmax flips in the output distribution and reorganizing substantial portions of the candidate token distribution, sometimes overriding tokens with higher probability by up to 2.39×2.39\times. Figure 1

Figure 1: Top-1024 overlap between iter=1 and iter=4 hidden states across the layer stack and generated sequence, revealing stable positions and cascading basin shifts as dark vertical streaks.

Latent computation through iterative refinement allows the model to explore alternate semantic trajectories. At position zero, universal basin shift occurs: every question demonstrates radical hidden-state reorganization upon additional iteration. Later passes effect genuine computational work, shifting or reinforcing basin affiliations and sometimes correcting or regressing previously correct answers, mirroring task-dependent optimal compute depth observed in other recurrent architectures. Figure 2

Figure 2: Layer profile of top-1024 overlap at low-overlap positions, showing non-monotonic divergence and localized activity peaks at specific layer bands.

Learned Blend Coefficient Structure

Training fully exploits per-dimension and per-layer freedom in the blend coefficients. Rather than shifting magnitudes globally, each layer develops distinct per-dimension adaptation patterns organized into stack-wide coherent trajectories (as demonstrated by principal component analysis projections). This confirms that blend parameterization is used for content-sensitive, layer-specific adaptation rather than simple bias replication. Figure 3

Figure 3

Figure 3: Per-layer alpha distributions, highlighting the adaptation and widening of percentile bands post-training.

Quantitative Evaluation and Numerical Results

SST V2 was co-trained into a 27B Gemma 3 backbone on CodeACT-reformulated GSM8K data. In staged compute evaluation, SST V2 achieved a +15.15+15.15 percentage point gain in accuracy over a matched fine-tuned baseline on the out-of-distribution GPQA-Diamond benchmark and cut baseline GSM8K errors by 46%46\%. On GPQA-Diamond, the 27B SST outperformed several open-weight and proprietary systems, including models up to 25×25\times larger.

Strong numerical findings:

  • GPQA-Diamond: SST V2 staged compute accuracy 61.11%61.11\% (baseline 45.96%45.96\%).
  • GSM8K: SST V2 staged compute kk0 (baseline kk1).
  • Error correction: SST V2 recovers kk2 of baseline errors.
  • Model scale: SST V2 surpasses DeepSeek V3 (671B), Gemini 2.0 Flash, and Qwen 2.5 72B in absolute accuracy at 27B parameters, using zero-shot greedy decoding.

Adaptive Halting and Overthinking Mitigation

A uniform iteration depth produces substantial regression on questions that do not require additional refinement: e.g., flat iter=4 accuracy on GPQA-Diamond drops to kk3, losing kk4pp against staged compute. The paper introduces an adaptive halting probe—a small MLP reading 107 essential hidden-state dimensions at layer 15 of position 0—that predicts whether further iteration will regress an answer. This probe closes kk5 of the iter=1-to-staged compute gap without breaking any correct answers, confirming the existence of generalizing features in the latent state that correlate with future computational trajectory. Figure 4

Figure 4: The 107 essential hidden-state dimensions at layer 15; only kk6 of the hidden state signal is necessary for halt prediction.

Figure 5

Figure 5: Layer 15 halt signal probe output, showing clear separation between must halt and safe timesteps, with zero overthinks.

Implications and Future Directions

SST V2 demonstrates that reasoning improvements are attributable to architectural mechanisms—not scale or dataset—establishing a new axis of computation orthogonal to parameter scaling and chain-of-thought. Practical implications include model deployment strategies leveraging adaptive iteration, robust error correction in reasoning tasks, and interpretability advances through direct tracing of semantic basin shifts in latent computation. The architecture is broadly compatible with any Lipschitz-bounded gated FFN transformer, suggesting potential for transferable reasoning improvements.

Theoretical implications:

  • SST V2 offers a framework for structured, emergent reasoning trajectories in continuous latent space, visible through basin dynamics and convergence across multiple computational paths.
  • Latent state features provide actionable signals for adaptive halting and future prediction, enabling new forms of reasoning competence and meta-cognition within LLMs.

Future developments:

  • Scaling SST V2 to larger models, RL finetuning with native halting, and task-specific adaptation of probe features.
  • Deeper iteration and cross-domain reasoning capability across diverse benchmarks.
  • Further mechanistic interpretability research on emergent structure in latent state dynamics.

Conclusion

State Stream Transformer V2 provides a parameter-efficient, parallel-trainable architecture for nonlinear latent space recurrence, enabling both horizontal and vertical axes of latent computation in transformers. Its two-pass training method unlocks practical co-training with existing backbones. SST V2 organizes latent space into stable, content-dependent basins, supports adaptive deliberation, and achieves competitive performance with much larger models on reasoning benchmarks. Its latent computation yields structured, causally prior basin transitions and emergent features enabling adaptive reasoning management, offering concrete utility for future developments in transformer-based architectures and interpretability research.

Paper to Video (Beta)

No one has generated a video about this paper yet.

Whiteboard

No one has generated a whiteboard explanation for this paper yet.

Open Problems

We haven't generated a list of open problems mentioned in this paper yet.

Collections

Sign up for free to add this paper to one or more collections.

Tweets

Sign up for free to view the 1 tweet with 2 likes about this paper.