StateLoop: Persistent State for Looped Transformers

Feichen Song; Boyi Zeng; He Li; Shixiang Song; Yiqin Hao; Zitong Wang; Yizhou Yang; Yifan Liu

cs.CLOct 26, 2026

StateLoop: Persistent State for Looped Transformers

Authors: Feichen Song, Boyi Zeng, He Li, Shixiang Song, Yiqin Hao, Zitong Wang, Yizhou Yang, Yifan Liu,

Organizations: Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute

Abstract

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.

Figures & tables

Figure 1: StateLoop architecture. (a) The looped model with a persistent state per token. (b) State update, optional token mixing, and readout.

W_{s\to z}

feeds the next iteration, whereas

W_{s\to out}

supplies the final prediction.

	Parameters		Perplexity $\downarrow$
Model	Params (M)	vs. Vanilla Looped	Pile	LAMBADA OAI	WikiText-2
Vanilla Looped	1420.92	0.00%	7.79	6.94	16.50
Vanilla Looped + FFN	1467.22	+3.26%	7.75	7.00	16.34
Hyperloop	1421.32	+0.03%	7.77	6.84	16.66
MeSH	1421.51	+0.04%	7.74	6.77	16.43
PonderLM	1420.92	0.00%	7.85	7.31	16.72

Table 1: Language-modeling PPL at 1.4B (30B training tokens). Full-backbone models use

K=4

, with MeSH and Parcae matched at 96 layer applications. LAMBADA-OpenAI uses zero-shot token PPL and WikiText-2 uses word PPL.

Zero-shot
Model	LAMBADA OAI	ARC-E	ARC-C	WinoG.	PIQA	Hella.	SciQ	RACE	Avg-8
Vanilla Looped	58.24	52.57	28.16	54.85	70.02	50.39	77.90	34.35	53.31
Vanilla Looped + FFN	59.05	50.93	26.71	56.91	69.42	50.50	75.60	33.68	52.85
Hyperloop	61.28	51.05	27.56	55.72	70.46	49.69	80.70	34.16	53.83
MeSH	60.94	50.88	26.62	59.67	69.70	50.52	78.10	34.55	53.87
PonderLM	58.26	50.59	28.58	55.41	68.93	48.96	78.00	33.68	52.80

Table 2: Zero-shot and five-shot accuracy (%) at 1.4B after 30B training tokens. Full-backbone models use

K=4

; MeSH and Parcae match their 96 layer applications. Avg-8 is the unweighted eight-task mean. LAMBADA-OAI denotes LAMBADA-OpenAI.

Figure 2: Loop depth and Pile PPL at 410M. Models are trained for 25B tokens at each plotted

K

. StateLoop uses identity mixing and

d=d_{s}=1024

. The PPL advantage over Vanilla Looped grows with trained loop depth.

Model	Params (M)	$K$	Latency (ms/token) $\downarrow$	Throughput (tok/s) $\uparrow$	Latency change
Vanilla Looped	1420.92	4	121.224 $\pm$ 1.099	8.249	baseline
Vanilla Looped	1420.92	5	151.527 $\pm$ 1.064	6.599	+24.997%
Vanilla Looped + FFN	1467.22	4	121.167 $\pm$ 1.293	8.253	$-0.047\%$
StateLoop	1437.71	4	122.943 $\pm$ 1.483	8.134	+1.418%
StateLoop + Attn	1454.49	4	123.778 $\pm$ 1.348	8.079	+2.106%
StateLoop + Mamba-2	1467.23	4	123.397 $\pm$ 0.923	8.104	+1.792%

Table 3: Decoding efficiency at 1.4B. Measurements use BF16, batch size 1, and a 1,024-token prefill followed by a 1,024-token continuation on one GPU. Latency is mean

\pm

standard deviation over five runs; changes are relative to Vanilla Looped at

K=4

Figure 3: State width at 410M (

K=7

, 25B tokens). Pile PPL versus parameter overhead for identity mixing and backbone width

d=1024

. Points are labeled with state width

d_{s}

and PPL; overhead is relative to the 406.90M-parameter Vanilla Looped baseline.

Ablation	Variant	Params (M)	PPL $\downarrow$
Feedback	State only	409.00	9.44
Feedback	StateLoop	411.10	9.04
Persistence	Reset state	411.10	9.25
Persistence	Persistent state	411.10	9.04
Update rule	Vanilla Looped	406.90	9.34
	StateLoop (GRU)	416.34	9.09

Table 4: Persistent-state ablations on Pile at 410M. All models use

K=7

and 25B training tokens. State variants use

d_{s}=1024

and identity mixing. State-only feedback uses the recurrent state as the next backbone input, whereas reset state sets the previous state to zero before each update. The GRU variant replaces the selective-SSM update while keeping feedback and final readout unchanged. Parameter counts are in millions.

Figure 4: State analysis at 410M (

K=7

, identity mixing). Evaluated on the full Pile validation set. (a) Shapley shares of the NLL increase from masking reads. (b) Calibrated single replacement versus frozen-state reuse; the dashed line marks the unmodified model.

Persistent Memory for Language Models

Persistent memory for language models studies how models retain, update, and use information across interactions and beyond their context windows.

See all 271 papers

LoopFold: Parallelizing Recurrent Depth in Looped Transformers
Oct 26, 2026Feichen Song, Boyi Zeng, Yiqin Hao +7Recurrent Transformers Efficient Transformer Inference
Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute
Asymmetric Looped Transformers: Fixed Global Representations for an Iterative Decoder
Oct 20, 2026Cheng Liang, Tengfei Zhang, Wenjie DU +6Recurrent Language Models Recurrent Transformers
Shanghai Jiaotong University · University of Science and Technology of China · Nanyang Technological University +1
FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates
Sep 24, 2026Wanqi Yang, Shiwei LiuEfficient Transformer Inference Recurrent Transformers
ELLIS Institute Tübingen Max Planck Institute for Intelligent Systems Tübingen AI Center

Oct 26, 2026J/K move · Enter open · S save

LoopFold: Parallelizing Recurrent Depth in Looped Transformers

Feichen Song, Boyi Zeng, Yiqin Hao, Shixiang Song+6

Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute

Recurrent Transformers Efficient Transformer Inference LLM Inference Acceleration

Open paper AskPDF arXiv

Looped Transformers increase computation per token by repeatedly applying a shared backbone, offering a way to improve language modeling and downstream task performance at a fixed backbone parameter count. However, each loop depends on the preceding loop, so increasing recurrent depth also increases decoding latency. Existing approaches include Parallel Loop Transformers, which reduce latency by overlapping loop computation across token positions, but yield rapidly diminishing marginal accuracy gains at larger depths. To address this challenge, we introduce LoopFold, which folds computation at multiple recurrent depths into parallel batches for the same token prefix. A lightweight predictor estimates inputs to later loops so they can start before earlier loops finish. A corrector then refines the outputs of these later loops using newly available results from earlier loops. All components are trained jointly using only the final next-token prediction loss. In experiments with 1.4B models, deeper LoopFold configurations achieve progressively higher average zero-shot accuracy across eight tasks, rising from 53.41% to 56.38%. Vanilla looped models and PLT show limited gains over the same depth range. Compared with a nine-loop vanilla model, LoopFold uses five sequential stages and additional parallel computation to reduce decoding latency by 46%.

Abstract

Persistent Memory for Language Models

Explore similar work

LoopFold: Parallelizing Recurrent Depth in Looped Transformers

Asymmetric Looped Transformers: Fixed Global Representations for an Iterative Decoder

FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates