PaperlayerView exact revision on arXiv

Source-linked AI summary

StateLoop: Persistent State for Looped Transformers

Feichen Song, Boyi Zeng, He Li, Shixiang Song, Yiqin Hao, Zitong Wang, Yizhou Yang, Yifan Liu, Ziwei He, Zhouhan Lin

arXiv:2610.35544v1cs.CL

TL;DR

Vanilla Looped relies on one hidden state both for the next backbone input and for information that may be useful several iterations later. StateLoop adds a gated persistent state per token, and reports lower Pile perplexity and higher average accuracy across eight zero-shot tasks than Vanilla Looped.

  • Problem

    Vanilla Looped feeds each backbone output directly into the next iteration, so retaining information useful several iterations later depends on the same hidden state used for subsequent computation.

    Source: §1 · p. 1

  • Method

    StateLoop maintains a gated persistent state per token, updating it with backbone output and projecting it back into the backbone stream for later iterations and final prediction.

    Source: §1 · p. 1

  • Results

    StateLoop reduces Pile perplexity by 0.18 and improves average accuracy across eight zero-shot tasks by 0.97 percentage points at 1.4B versus Vanilla Looped.

    Source: §1 · pp. 1–2

  • Takeaways & Limitations

    StateLoop’s Pile-perplexity advantage over Vanilla Looped grows from approximately 0.14 at K=3 to 0.30 at K=7, consistent with persistent state becoming increasingly useful as recurrent depth increases.

    Source: §4.3 · p. 6

    Experiments cover models up to 1.4B, and the controlled loop-depth analysis is conducted at 410M.

    Source: §6 · p. 11

Abstract

from arXiv · show

Looped Transformers offer a parameter-efficient approach to language modeling by repeatedly applying a shared backbone. Vanilla looped Transformers use each backbone output directly as the input to the next iteration. These outputs may contain information useful several iterations later, but retaining it across the loop depends entirely on the single hidden state. We introduce StateLoop, which maintains one additional persistent state per token to retain information across iterations. At each iteration, a learned gate controls how much of the previous state is retained and how much information from the backbone output is incorporated. The updated state is projected back and added to the backbone output for subsequent iterations and final prediction. This design reduces Pile perplexity and improves average accuracy across eight zero-shot tasks over Vanilla Looped. Its perplexity advantage over Vanilla Looped grows with trained loop depth, consistent with persistent state becoming increasingly useful as recurrent computation deepens.

1. Introduction

Vanilla Looped uses one hidden state both to feed later iterations and to retain information that may be useful several iterations later. StateLoop adds a gated persistent state per token, and it improves language-modeling and downstream results over Vanilla Looped.

1. Introduction

  • Vanilla Looped uses each backbone output as the next iteration’s input, leaving one hidden state to carry information that may be useful several iterations later.The motivation draws on DenseFormer’s finding that access to earlier outputs can lower perplexity than access restricted to recent outputs (Pagliardini et al., 2024).Source: §1 · p. 1
  • StateLoop maintains an additional persistent state per token, gated to retain its previous value and incorporate a projection of the backbone output.The updated state is projected back and added to the backbone output for later iterations and final prediction.Source: §1 · p. 1
  • 0.18 lower Pile perplexity and 0.97 percentage points higher average accuracy across eight zero-shot tasks at 1.4B are reported for StateLoop versus Vanilla Looped.StateLoop’s advantage in perplexity also grows with trained loop depth.Source: §1 · pp. 1–2
  • StateLoop’s perplexity advantage over Vanilla Looped grows with trained loop depth, consistent with persistent state becoming increasingly useful as recurrent computation deepens.Resetting the state at every iteration increases perplexity.Source: §1 · p. 2

2. Related Work

Prior work develops looped Transformers for recurrent computation and language modeling, while varying how information is carried between iterations. StateLoop’s related-work positioning is a single continuously updated state per token alongside a direct recurrent hidden stream.

Looped Transformers and recurrent depth

  • The Universal Transformer introduced recurrent computation across depth; later work studies iterative algorithms for in-context learning (Yang et al., 2024) and large-scale pretraining with learned depth allocation (Zhu et al., 2025).Ouro extends the approach to large-scale language-model pretraining.Source: §2 · p. 2

Connections between iterations

  • Parcae, PonderLM, and Hyperloop modify or enrich the recurrent stream, whereas StateLoop retains the direct backbone stream and adds a separate persistent state (Prairie et al., 2026; Zeng et al., 2026; Zeitoun et al., 2026).Parcae constrains its transition matrix’s spectral norm, PonderLM feeds back a weighted sum of token embeddings, and Hyperloop carries multiple residual streams.Source: §2 · p. 2

Persistent state and memory

  • Memory-related approaches include static learned banks, dynamically updated memory tokens, and a multi-slot buffer; StateLoop instead uses one continuously updated state per token (Frey et al., 2026; Sapunov, 2026; Yu et al., 2026).StateLoop retains a direct main hidden stream and shares its state-update and feedback parameters across iterations.Source: §2 · p. 2
  • Recurrent state updates in LSTMs, GRUs, and selective state-space models and memory tokens passed between sequence segments provide related precedents (Cho et al., 2014; Gu and Dao, 2024; Bulatov et al., 2022).These approaches update state along a sequence or carry memory between segments.Source: §2 · p. 2

3. Method

StateLoop adds a persistent state per token to a looped Transformer, updating it from backbone outputs and reading it back into the backbone stream. The method includes selective-SSM updates, optional token mixing, and a final readout trained with autoregressive loss.

3. Method

  • StateLoop maintains a persistent state per token and reads the updated state into the next backbone input or final prediction.The architecture also permits optional token mixing before readout.Source: §3 · p. 2, Figure 1 · p. 3

3.1. Backbone Feedback and Persistent State

  • At each iteration, the state branch processes backbone outputs, updates and normalizes the state, then projects its readout back into the backbone-width input; parameters are shared across iterations.The state width may differ from the backbone width, and the state read is added to the backbone output for the next iteration.Source: §3.1 · p. 3

3.2.1. Selective SSM Update

  • The selective-SSM update uses an input-dependent step size to set gates that weight the previous state and projected backbone output.The learned retention and write weights can vary by token and iteration; the update recurrence operates along the loop axis independently for each token.Source: §3.2.1 · p. 4

3.2.2. Optional Token-Axis Mixing

  • Optional causal-attention or Mamba-2 mixers exchange information across token positions, while identity mixing leaves the updated state unmixed.The attention mixer uses pre-normalization and a residual connection; the Mamba-2 mixer applies a residual block along the token axis.Source: §3.2.2 · p. 4

3.3. Final Prediction and Training

  • At the final iteration, a separate readout projection combines the normalized state with the backbone output, and the model is trained end-to-end with standard autoregressive loss.The language-model head maps the resulting representation to vocabulary logits.Source: §3.3 · pp. 4–5

4. Experiments

The experiments compare StateLoop with looped and other baselines on language modeling, downstream accuracy, loop depth, and decoding efficiency. StateLoop improves language-modeling and downstream results over Vanilla Looped, with its Pile perplexity advantage growing at greater trained loop depths and a latency advantage over a fifth iteration.

4.1. Experimental Setup

  • The 1.4B comparisons use a shared 30B-token training setup, while the 410M studies use 25B tokens; baselines are implemented in the same codebase.The shared 1.4B setup also fixes the tokenizer, data pipeline, batch size, sequence length, schedule, and backbone; full-backbone recurrent models use four iterations.Source: §4.1 · p. 5
  • Evaluation reports PPL on Pile, LAMBADA-OpenAI, and WikiText-2, plus accuracy on eight downstream tasks under zero-shot and five-shot settings.Avg-8 is the unweighted average across the eight tasks.Source: §4.1 · pp. 5–6

4.2. Main Results

  • StateLoop lowers perplexity on all three datasets and improves average downstream accuracy under both prompting settings; both token mixers further improve Avg-8.Tables 1 and 2 cover language-modeling PPL and zero-shot and five-shot accuracy, respectively.Source: §4.2 · p. 6, Table 1 · p. 6, Table 2 · p. 7
  • Allocating additional parameters to the state branch outperforms widening the FFN; StateLoop does so with fewer parameters, while StateLoop + Mamba-2 Mixer does so at comparable parameter count.The comparison covers Pile PPL and both downstream averages.Source: §4.2 · p. 6

4.3. Loop Depth

  • At fixed state width, the Pile PPL advantage over Vanilla Looped grows from approximately 0.14 at K=3 to 0.30 at K=7.The 410M paired models use identity mixing, 25B training tokens, and the depth used during training; Figure 2 shows the depth comparison.Source: §4.3 · p. 6, Figure 2 · p. 7

4.4. Decoding Efficiency

  • At K=4, the state branch adds 1.42%–2.11% latency, versus 25.00% for a fifth Vanilla Looped iteration; StateLoop is 18.86% lower latency than that K=5 baseline.The comparison measures decoding efficiency for the 1.4B models.Source: §4.4 · pp. 6–7, Table 3 · p. 8

5. Ablation and Analysis

At 410M and seven iterations, the ablations vary state width and persistent-state design, while interventions analyze state reads and continued updates. The results show benefits from compact state, cross-iteration persistence, and later state updates.

5.1. State Width

  • Halving state width preserves most of the PPL gain while roughly halving added parameter cost, and even the narrowest state improves over Vanilla Looped.Figure 3 plots Pile PPL against parameter overhead for state widths at 410M, with overhead relative to the 406.90M-parameter Vanilla Looped baseline.Source: Figure 3 · p. 8, §5.1 · p. 8

5.2. Persistent-State Design

  • StateLoop achieves lower PPL than state-only feedback, favoring a direct connection between iterations plus a state read over replacing that path with recurrent state.The ablation compares these feedback designs in the persistent-state study.Source: §5.2 · p. 9, Table 4 · p. 9
  • 9.25 vs. 9.34 PPL shows reset state still beats Vanilla Looped; persistent state reaches 9.04 PPL under the same parameter count and training setup.The comparison isolates an additional benefit from carrying state across iterations.Source: §5.2 · p. 9
  • The GRU update preserves most of the improvement over Vanilla Looped, while the selective-SSM update achieves lower PPL with fewer parameters.The comparison keeps feedback, final readout, and identity mixing unchanged.Source: §5.2 · pp. 8–9

5.3. Analysis of the Persistent State

  • 63.70% of attributed NLL increase comes from feedback reads and 36.30% from the terminal read, with contributions concentrated in R5, R6, and R7.The analysis evaluates all 27 combinations of masked reads using exact Shapley values.Source: Figure 4 · p. 10, §5.3 · p. 10
  • 9.08 PPL follows replacing S3 once at S4, while freezing S3 across later boundaries raises PPL to 18.41.Both interventions continue running the backbone; the contrast indicates that earlier states remain partly usable nearby while continued updates matter as computation proceeds.Source: Figure 4 · p. 10, §5.3 · p. 10

6. Conclusion

StateLoop adds one persistent state per token to retain information across iterations, reducing perplexity and improving downstream accuracy over Vanilla Looped. The experiments reach 1.4B scale, while controlled loop-depth analysis is limited to 410M.

6. Conclusion

  • StateLoop reduces perplexity and improves average downstream accuracy over Vanilla Looped; causal-attention and Mamba-2 state mixers provide further downstream gains.Gains with both selective-SSM and GRU updates support the persistent-state design.Source: §6 · p. 11
  • StateLoop’s perplexity advantage over Vanilla Looped grows at greater trained depths, highlighting the value of persistent state for deeper recurrent computation.State interventions also indicate that earlier states remain partly usable at later boundaries while continued updates matter for prediction.Source: §6 · p. 11
  • Experiments cover models up to 1.4B, but controlled loop-depth analysis is conducted at 410M.The paper identifies larger-model analysis as a way to further characterize the interaction between persistent state and recurrent depth.Source: §6 · p. 11

Related papers

Loading 2610.35544v1…