The Surprising Effectiveness of Shared Memory in Looped Transformers

Giovanni Monea; Keshav Ramji; Yousef El-Kurdi; Luis A. Lastras; Yoav Artzi; Nathan Godey; Ramón Fernandez Astudillo

cs.LGOct 1, 2026

The Surprising Effectiveness of Shared Memory in Looped Transformers

Authors: Giovanni Monea, Keshav Ramji, Yousef El-Kurdi, Luis A. Lastras, Yoav Artzi, Nathan Godey, Ramón Fernandez Astudillo

Organizations: Cornell University · IBM Research

Abstract

Looped Transformers apply the same layers several times per token, adding compute to improve quality without more parameters. Each recursion, however, writes its own key-value cache, so memory still grows with compute. Inference-time techniques can shrink this cache at a cost in quality. We pretrain looped language models to share memory: only the first recursion writes a cache, and later recursions read it while keeping a short window of their own. Surprisingly, we find that sharing memory does not cost quality and instead improves it. At 150M-1B parameters, our Looped Prediction Transformer (LPT) and its hybrid variant set a new quality-memory frontier for looped models: with five recursions, the hybrid lowers validation perplexity on FineWeb-Edu by 1.12-1.82 relative to a same-size standard Transformer while using 76-79% less context memory. Through an extensive analysis, we investigate why memory sharing helps. Shared and local memory develop different representations, and later recursions attend mostly to the shared memory, which also acts as a gradient highway to the first recursion.

Figures & tables

Panel of Figure 1 — Figure 1: Shared memory not only saves memory, it also improves quality. Left: in our Looped Prediction Transformer ( lpt ), later recursions read the KV cache written by the first recursion and keep only a short window of their own. Right: perplexity and context memory of $1$ B-parameter models with two to five recursions, against a parameter-matched standard Transformer ( std ). Each of our models adds shared memory to the looped baseline above it: lpt to the looped Transformer lt , and h-lpt to the gated DeltaNet hybrid lt2 .

		$R=1$		$R=2$				$R=3$				$R=5$
Size	Model	PPL $\downarrow$	Acc. $\uparrow$	PPL $\downarrow$	Acc. $\uparrow$	KV $\downarrow$	FLOPs $\downarrow$	PPL $\downarrow$	Acc. $\uparrow$	KV $\downarrow$	FLOPs $\downarrow$	PPL $\downarrow$	Acc. $\uparrow$	KV $\downarrow$	FLOPs $\downarrow$
	std	$16.83$	$51.5$	$17.07$	$50.2$	$1.50$	1.19	$17.27$	$50.7$	$1.88$	1.31	$17.69$	$50.4$	$2.50$	1.54
	loopie			$16.58$	$51.6$	$2.00$	$1.83$	$16.50$	$51.9$	$3.00$	$2.66$	$16.56$	$52.1$	$5.00$	$4.32$
	lt			$16.17$	$51.5$	$2.00$	$1.83$	$15.98$	$52.1$	$3.00$	$2.66$	$16.00$	$52.4$	$5.00$	$4.32$
	plt			$15.84$	52.7	$1.02$	$1.84$	$15.55$	$52.9$	$1.03$	$2.67$	$15.49$	$53.5$	$1.06$	$4.34$
	lt2			15.76	$51.9$	0.35	1.29	$15.47$	$52.5$	0.52	1.86	$15.16$	$53.7$	0.87	2.98

Table 1: lpt and h-lpt reach lower perplexity than their looped baselines with much smaller caches. Main results across S, M, and L: validation perplexity, mean five-task accuracy (%), and context memory and FLOPs relative to same-size std . In the std rows,

R=1

is the standard Transformer, and

R>1

is deepstd . Elsewhere,

R

counts model passes. Bold and underline mark the best and second-best within each size and column group; our rows are shaded. We replicate the S results with three additional seeds and find the gains of Memory Sharing statistically significant (

p<0.01

, Appendix B ).

Figure 2: lpt and h-lpt reach higher accuracy with much smaller caches than the looped baselines. Mean five-task accuracy vs. context memory at each size. Curves follow two, three, and five recursions, with larger markers at five. Dashed lines mark single-pass std .

Panel of Figure 3 — Figure 3: lpt trains best with final loss, while lt consistently worsens with more recursions without adaptive loss. Validation perplexity ( $\downarrow$ ), collapsed runs in the upper strip on a log scale. Filled markers are final loss, hollow ones adaptive loss.

Panel of Figure 5 — Figure 5: Prediction recursions attend mostly to shared memory. Attention mass assigned to shared KV entries in lpt , averaged over queries and heads. For clarity, only six layers are displayed. The bottom row averages all layers.

Panel of Figure 7 — Figure 7: Shared memory restores much of the gradient reaching the first recursion. Mean gradient norm contributed by each recursion, with single-pass std as reference. The first lpt bar is split into shared-KV (hatched) and hidden-state projections. lt is trained with adaptive loss, the others with final loss. Whiskers are $95\%$ sequence-bootstrap intervals.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

	S	M	L
Layers ( $L$ )	12	24	36
Hidden size ( $d$ )	768	1,024	1,280
Heads	12	16	20
FFN size	3,072	4,096	5,120
Parameters	151.9M	454.2M	1.01B

Appendix

Table 2: Base model configurations. Looped models reuse the listed

L

-layer stack.

Model	$R=1$	$R=2$	$R=3$	$R=5$
std	$16.84\pm 0.03$	$17.03\pm 0.05$	$17.27\pm 0.07$	$17.64\pm 0.13$
sps		$16.12\pm 0.07$
loopie		$16.57\pm 0.03$	$16.57\pm 0.08$	$16.58\pm 0.03$
lt		$16.16\pm 0.03$	$16.00\pm 0.05$	$15.99\pm 0.03$
plt		$15.83\pm 0.05$	$15.55\pm 0.03$	$15.46\pm 0.03$
lt2		$15.77\pm 0.02$	$15.47\pm 0.02$	$15.16\pm 0.01$

Appendix

Table 3: Seed variance at S. Final validation perplexity, mean

\pm

95\%

confidence interval (Student-

t

) over four seeds. In the std row,

R=1

is the standard Transformer and

R>1

the deeper, narrower parameter-matched Transformer. sps is the two-recursion baseline of Appendix E .

Comparison	$R$	Gap (PPL)	Gap $/$ SD	$p$
lpt vs. lt	2	$-0.40\pm 0.05$	$13$	$6.5\times 10^{-5}$
	3	$-0.71\pm 0.05$	$24$	$9.4\times 10^{-6}$
	5	$-1.05\pm 0.04$	$39$	$2.4\times 10^{-6}$
h-lpt vs. lt2	2	$-0.04\pm 0.01$	$4$	$1.7\times 10^{-3}$
	3	$-0.07\pm 0.02$	$5$	$9.8\times 10^{-4}$
	5	$-0.13\pm 0.02$	$9$	$1.6\times 10^{-4}$

Appendix

Table 4: Paired tests at S. Same-seed perplexity gap of the first model over the second, mean

\pm

95\%

confidence interval over four seeds, the gap in units of the standard deviation of the same-seed differences, and the one-sided paired

t

-test for the first model having lower perplexity.

			S			M
Model	Normalization	Loss	$R=2$	$R=3$	$R=5$	$R=2$	$R=3$	$R=5$
lt	Ouro (main)	Adaptive	16.17	15.98	16.00	13.28	13.13	13.34
	Ouro	Final	16.05	16.32	2,040	17.10	31.74	2,090
	Pre-norm only	Final	16.28	15.99	15.77	13.28	13.05	12.83
	Sandwich only	Final	16.23	15.83	15.67	–	–	–
	Input norm	Final	16.10	15.79	15.49	13.05	12.88	13.02

Appendix

Table 5: lpt outperforms lt with the best normalization we tried. Validation perplexity (

\downarrow

) after

20

B tokens. Bold marks the best per column and underline the best lt . Values above

2{,}000

are unigram-level. A dash marks a run that did not finish.

Figure 9: Shared and local keys remain distinguishable with full prediction caches. Expanded comparison for Figure 6 , including $\omega=4{,}096$ .

Figure 10: Query/key geometry across lpt models. S and M models at three and five recursions, with final loss and $\omega=64$ . Columns match relative layer depth.

Figure 11: Attention to shared memory across training objectives and window sizes. Expanded comparison for Figure 6 . Mean rows average all layers.

Figure 12: Shared-cache feedback persists across training objectives and window sizes. Expanded version of Figure 8 , with the same sequences and final-loss gradient probe in every condition.

Figure 13: Shared-KV feedback relative to the hidden-state path. Gradient-norm ratios at fresh initialization and after training. Whiskers: $95\%$ sequence-bootstrap intervals.

	S			M
Model	$R=2$	$R=3$	$R=5$	$R=2$	$R=3$	$R=5$
lt (final loss)	16.05	16.32	2,040	17.10	31.74	2,090
lpt (detached KV)	16.34	16.14	33.21	13.87	1,995	1,527
lpt	15.73	15.28	14.91	12.91	12.55	12.33

Appendix

Table 6: Detaching shared-cache gradients hurts learning. Validation perplexity (

\downarrow

) at S and M, all models trained with final loss; values above

1{,}000

are runs that failed to train. Bold marks the best per column.

S	$R=1$		$R=2$				$R=3$				$R=5$				$R=9$
Model	PPL $\downarrow$	Acc. $\uparrow$	PPL $\downarrow$	Acc. $\uparrow$	KV $\downarrow$	FLOPs $\downarrow$	PPL $\downarrow$	Acc. $\uparrow$	KV $\downarrow$	FLOPs $\downarrow$	PPL $\downarrow$	Acc. $\uparrow$	KV $\downarrow$	FLOPs $\downarrow$	PPL $\downarrow$	Acc. $\uparrow$	KV $\downarrow$	FLOPs $\downarrow$
std	$16.83$	$51.5$	$17.07$	$50.2$	$1.50$	1.19	$17.27$	$50.7$	$1.88$	1.31	$17.69$	$50.4$	$2.50$	1.54	$18.32$	$49.1$	$3.38$	1.81
sps			$16.18$	52.6	$1.02$	$1.84$
loopie			$16.58$	$51.6$	$2.00$	$1.83$	$16.50$	$51.9$	$3.00$	$2.66$	$16.56$	$52.1$	$5.00$	$4.32$	$16.98$	$51.0$	$9.00$	$7.64$
lt			$16.17$	$51.5$	$2.00$	$1.83$	$15.98$	$52.1$	$3.00$	$2.66$	$16.00$	$52.4$	$5.00$	$4.32$	$16.13$	$52.2$	$9.00$	$7.64$
plt			$15.84$	52.7	$1.02$	$1.84$	$15.55$	$52.9$	$1.03$	$2.67$	$15.49$	$53.5$	$1.06$	$4.34$	$15.51$	$53.2$	1.12	$7.68$

Appendix

Table 7: Quality and efficiency through nine recursions. One table per size. Metrics and formatting follow Table 1 : in the std row,

R=1

is the standard Transformer and

R>1

is deepstd . sps appears in the

R=2

column, whose memory and FLOPs it matches.

Figure 14: Quality and memory through nine recursions. Validation perplexity, mean five-task accuracy, and context memory relative to std .

Figure 15: Quality versus context memory across model sizes. Curves vary recursion count. Marker size denotes recursions, and shade denotes model size. S/M train for

20

B tokens and L for

50

Figure 16: Quality and memory across model sizes. Mean five-task accuracy and context memory for std and looped models at nine recursions.

Figure 17: h-lpt is better than std at similar or higher throughput. Validation perplexity (top) and mean five-task accuracy (bottom) against generated tokens per second at each model’s maximum batch, relative to same-size std , on one H100 at a

4{,}096

-token context. Curves follow two, three, and five recursions, with larger markers at five. Dashed lines mark std .

Memory-Efficient Inference

Memory-Efficient Inference reduces the memory needed to run neural networks while balancing accuracy, latency, throughput, and energy use.

See all 164 papers

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
Adaptive Loops and Memory in Transformers: Think Harder or Know More?
Mar 9, 2026Markus Frey, Behzad Shomali, Ali Hamza Bashir +3Transformer Memory-Augmented Language Models
Lamarr Institute · Fraunhofer IAIS · University of Bonn
Sparse Layers are Critical to Scaling Looped Language Models
May 9, 2026Ryan Lee, Jacob Biloki, Edward J. Hu +1Mixture-of-Experts Language Models Recurrent Transformers
1USC Information Sciences Institute · 2Netflix · 3Independent Researcher

Sep 24, 2026J/K move · Enter open · S save

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

Wanqi Yang, Shiwei Liu

ELLIS Institute Tübingen Max Planck Institute for Intelligent Systems Tübingen AI Center

Efficient Transformer Inference Recurrent Transformers

Open paper AskPDF arXiv

Looped Transformers have attracted substantial attention as a parameter-efficient approach to increasing computational depth through repeated application of shared Transformer blocks. However, their practical advantages over conventional Transformers remain under debate: each additional loop incurs another Transformer pass and requires caching another set of KV states, causing inference FLOPs and KV-cache memory to grow continuously with loop depth. This overhead becomes particularly severe at large loop counts and long context, preventing the parameter efficiency of Looped Transformers from translating into practical inference efficiency. In this paper, we find that much of the additional computation and storage introduced by looping is redundant. As recurrence proceeds, state changes become increasingly concentrated on a small subset of tokens; attention-output differences are dominated by a sparse and stable subset of key columns; and KV residuals between adjacent loops become progressively more amenable to low-bit quantization. Building on these observations, we introduce FlashLoop, a training-free inference framework that reduces cross-loop redundancy through token-sparse updates, sparse attention, and KV-residual quantization. Across several Looped Transformers models, FlashLoop delivers lossless accuracy while achieving up to 1.64

\times

end-to-end speedup and up to 6

\times

KV-cache memory reduction, substantially improving the practicality of scaling Looped Transformers to greater computational depths and longer context.

Abstract

Memory-Efficient Inference

Explore similar work

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

Adaptive Loops and Memory in Transformers: Think Harder or Know More?

Sparse Layers are Critical to Scaling Looped Language Models