PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.02383v1cs.LGcs.AI

TL;DR

Looped Transformers add computation through repeated layer application, but their recursion-specific caches make context memory grow with depth. This work trains looped models to share memory, and finds that LPT and H-LPT improve quality while reducing cache use; at five recursions, H-LPT surpasses a same-size standard Transformer with substantially less context memory.

  • Problem

    Looped Transformers keep a separate key-value cache for each recursion, so context memory grows with depth and can exceed weight memory at long contexts.

    Source: §1 · p. 2

  • Method

    The first recursion writes full-context memory, which later recursions read alongside short local windows, and the models are pretrained with this sharing design.

    Source: §3 · p. 3

  • Results

    At five recursions, H-LPT lowers validation perplexity by 1.12–1.82 relative to a same-size standard Transformer while using 76–79% less context memory.

    Source: §5 · p. 6

  • Takeaways & Limitations

    The analyses suggest shared and local memory support different recursion roles, while shared-memory reads provide a direct gradient path to the first recursion.

    Source: §8 · p. 10

    The study trains models up to 1B parameters on one corpus with 4,096-token contexts and does not test long-context tasks or very advanced tasks.

    Source: §8 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Looped Transformers add computation through repeated layer application but retain separate recursion-specific caches, making long-context memory a limiting cost. This paper trains looped models to share memory and finds that sharing improves quality while reducing cache use.

1 INTRODUCTION

  • Memory sharing is trained as an architectural inductive bias, prompted by evidence that untrained recursions can reuse one another’s caches.The earlier reuse work studied sharing mainly as a memory-saving technique.Source: §1 · p. 2
  • LPT improves over full-cache LT even when each later recursion’s own cache covers less than 2% of the context.Across sizes, it lowers validation perplexity by up to 1.09 nats and raises downstream accuracy by up to 4.8 percentage points, with up to 4.7× less context memory and nearly the same inference FLOPs.Source: §1 · p. 2
  • Later recursions rely heavily on shared memory, whose reads provide a direct gradient path to the first recursion and bypass intermediate recursions.Shared and local memory develop distinct key representations; blocking the shared-memory gradient path leaves quality and stability below LPT.Source: §1 · p. 2

2 LOOPED TRANSFORMERS

Looped Transformers reuse one stack of layers across successive recursions, but each recursion computes its own hidden states and key-value entries. Consequently, full-context cache memory grows with recursion depth even though the parameter count stays fixed.

2 LOOPED TRANSFORMERS

  • Each recursion applies the same shared stack of L Transformer blocks, giving effective depth L · R with the parameters of one L-layer stack.The formulation uses the full-stack architecture described as Ouro’s formulation, building on recurrent-depth architectures.Source: §2 · pp. 2–3
  • Each layer computes keys and values from the hidden states at its current recursion, and vanilla LT reads only that recursion’s cache.A query at position i reads the causal prefix of its own cache.Source: §2 · p. 3
  • At inference, LT retains R full-context caches per layer, so context memory grows linearly with recursion depth despite fixed parameter count.Caches may be dropped or merged during decoding, but the model is trained with a separate cache for every recursion.Source: §2 · p. 3
  • Hybrid looped models replace most full-attention layers with linear-attention layers whose recurrent state does not grow with context length.LT2 uses Gated DeltaNet layers and one full-attention layer per five, but each remaining full-attention layer still has a separate cache at every recursion.Source: §2 · p. 3

3 LOOPING WITH SHARED MEMORY

Memory Sharing makes the first recursion the source of full-context memory and lets later recursions combine that shared cache with short local windows. The design extends to hybrid looped Transformers, retaining their own recurrent states in linear-attention layers.

3 LOOPING WITH SHARED MEMORY

  • The first recursion writes a full-context cache; each later recursion attends to that shared cache and its own KV entries in one softmax.This shared-memory rule is applied at each full-attention layer.Source: §3 · p. 3
  • LPT uses the first recursion as a state recursion that builds full-context memory, while prediction recursions use it to refine predictions with short local caches.The model applies the same L-layer stack R times before passing the final output to the language-modeling head.Source: §3 · p. 3
  • With ω = T, prediction recursions have full-context caches of their own, isolating memory sharing from the effect of using a short window.This setting is used to separate the effects of sharing and window restriction.Source: §3 · p. 3
  • In hybrid looped models, later linear-attention layers keep their own small recurrent states, while full-attention layers apply the LPT sharing rule.Applying this to LT2 yields H-LPT.Source: §3 · p. 3
  • Table 1 reports validation perplexity, mean five-task accuracy, context memory, and FLOPs for the main results across S, M, and L sizes.Its caption notes that the S-size memory-sharing gains were replicated with three additional seeds and were statistically significant at p < 0.01.Source: Table 1 · p. 4
  • At five recursions, a 64-token window increases attended KV entries per generated token by only 1.25% while reducing full-attention cache relative to LT by 4.7× at a 4,096-token context.The calculation uses T = 4,096 and ω = 64.Source: §3 · p. 4

4 EXPERIMENTAL SETUP

The experiments compare memory-sharing models with looped and non-looped baselines across model sizes, recursion counts, and language-modeling and downstream evaluations. Training uses FineWeb-Edu, with additional seeds at the smallest size to check variation.

4 EXPERIMENTAL SETUP

  • The baselines include LT, LT2, PLT, STD, LOOPIE, and DEEPSTD, covering looped, parallel-loop, standard, and deeper non-weight-shared models.DEEPSTD matches looped models in effective depth and approximately in parameters.Source: §4 · p. 4
  • Figure 2 plots mean five-task accuracy against context memory at each size, with curves for two, three, and five recursions and dashed lines for single-pass STD.The caption describes LPT and H-LPT as reaching higher accuracy with smaller caches than looped baselines.Source: Figure 2 · p. 5
  • Models span approximately 150M, 450M, and 1B parameters, use 4,096-token contexts, and are tested with two, three, or five recursions.LPT and H-LPT use ω = 64 unless otherwise specified.Source: §4 · pp. 4–5
  • Models are pretrained on FineWeb-Edu for 20B tokens at S and M sizes and 50B at L, with matched initialization seeds and data order in the main comparisons.Three further seeds per model are trained at S to check seed variation.Source: §4 · p. 5
  • The study uses final loss or adaptive loss, with LT using adaptive loss and other models using final loss unless specified otherwise.Adaptive loss weights per-recursion NLLs using a learned per-token halting distribution and adds KL regularization.Source: §4 · p. 5
  • Evaluation measures held-out FineWeb-Edu perplexity, mean zero-shot accuracy on five tasks, and context memory and decoding FLOPs at 4,096 tokens.The five tasks are ARC-Easy, HellaSwag, PIQA, SciQ, and LAMBADA.Source: §4 · p. 5

5 RESULTS

Memory sharing improves quality over looped and standard Transformer baselines while substantially reducing context memory, especially in the hybrid model. Results are reported across recursion counts, with normalization and depth-related limitations noted.

Memory Sharing Improves Quality

  • At matched size and recursion count, LPT lowers validation perplexity by 0.38–1.09 and raises accuracy by 0.4–4.8 points over LT, with 2.0–4.7× less KV memory.LT stops improving beyond three recursions at every size, while LPT improves through five.Source: §5 · p. 5
  • Nine recursions bring instability in several model families and limited accuracy gains.This boundary is consistent with the reported loss spikes for Ouro at eight recursions.Source: §5 · p. 5
  • At five recursions, H-LPT lowers perplexity by 0.09–0.15 and raises accuracy by 0.8–1.2 points over LT2, with 4.1–4.2× less context memory.Across recursion counts, H-LPT matches or outperforms LT2 in quality while keeping context memory at 0.18–0.24× STD’s.Source: §5 · p. 6
  • At low recursion counts, LT2 improves on LT while using less memory and compute; alternative normalizations make LT more robust but none matches LPT.The authors note Ouro’s normalization as a possible factor in LT’s gap and instability.Source: §5 · p. 6

More Quality from the Same Parameters

  • At five recursions, H-LPT lowers perplexity by 1.12–1.82 and raises accuracy by 3.1–5.6 points versus STD, using 76–79% less memory.It fits 3.6–5.0× as many sequences as STD on one GPU.Source: §5 · p. 6
  • LPT and especially H-LPT improve the quality–memory frontier for looped models.Source: §5 · p. 6

6 WHY MEMORY SHARING HELPS

Memory Sharing stabilizes deeper LPT training and retains its advantage across prediction-window sizes. Analyses link this to shared-memory attention, distinct keys, and gradient feedback to the first recursion.

Memory Sharing Stabilizes Deep Recursion

  • LPT stays stable with final loss as recursion depth grows, while LT without adaptive loss degrades and collapses to unigram-level loss at five recursions.Final loss lowers LPT perplexity by 0.16–0.50 compared with adaptive loss.Source: §6 · p. 6, Figure 3 · p. 7

Prediction Recursion Needs Local Memory

  • Across ω = 0–4,096, LPT beats LT by 0.24–1.15 perplexity; at ω = 64, gains are 0.32–1.09, making window size's effect modest.LT uses its better objective per depth in this comparison; ω = 64 is the adopted default.Source: §6 · p. 6, Figure 4 · p. 7

Shared and Local Memory Develop Different Representations

  • Shared and local keys cluster separately in LPT, unlike keys across LT recursions; they remain distinguishable with full prediction caches, so similar representations are not required for sharing.The authors suggest this separation reflects memory-building and prediction roles.Source: §6 · p. 7, Figure 6 · p. 8

Later Recursions Attend More to Shared Memory

  • Shared memory receives 57–71% of attention in earlier prediction recursions and 73–80% in the final recursion.The figure averages attention to shared keys over queries and heads.Source: §6 · p. 7, Figure 5 · p. 8

Shared Memory Is a Gradient Highway

  • LPT restores much of the first-recursion gradient relative to STD as depth grows, mostly through shared KV at five recursions; detaching this path worsens perplexity and often breaks deep training.The detachment comparison keeps the forward computation and hidden-state gradient path unchanged.Source: §6 · p. 8

7 RELATED WORK

Prior work reuses caches to reduce memory or separates memory construction from prediction across streams or layers. LPT instead makes that split across loop recursions, and differs from parallel cache-sharing designs.

Separating Memory from Prediction

  • LPT splits memory construction and prediction across recursions, rather than token streams (Monea et al., 2026) or layers (Sun et al., 2024; 2026; DeepSeek-AI, 2026; Capps, 2026).The first recursion writes full-context memory; later recursions read it while predicting and keep a short local window.Source: §7 · pp. 8–9

Looped Transformers and Memory Sharing

  • Inference-time cache reuse includes Huginn and Ouro (Zhu et al., 2025), while trained variants reuse, precompute, gate, or compress caches.Keeping only the first recursion’s cache sharply reduces Ouro’s quality; the paper argues sharing needs training.Source: §7 · p. 9
  • Unlike PLT (Wu et al., 2025), which parallelizes recursions with gated local caches, LPT is sequential and reads both caches in one softmax; the authors report that LPT performs better.The paper frames sharing as an inductive bias that separates memory from prediction.Source: §7 · p. 9

8 DISCUSSION

The discussion concludes that memory sharing improves quality while reducing context memory, and proposes that shared and local memory support different roles. It also notes limits in compute efficiency, evaluation scope, and token-wise adaptivity.

8 DISCUSSION

  • Shared-memory LPT models improve perplexity and downstream accuracy over usual looped counterparts while matching a standard Transformer’s KV-cache memory.The analysis suggests the first recursion builds memory that later recursions use to refine predictions, while those reads provide a direct gradient path back to the first recursion.Source: §8 · p. 10
  • Shared and local memory develop different representations, supporting a division in which later recursions use the first recursion’s memory to refine predictions.Reading shared memory also gives the final loss a direct gradient path to the first recursion, and cutting this path hurts learning.Source: §8 · p. 10
  • The authors conclude that shared memory is a useful inductive bias for making better use of computation within a limited memory budget.Source: §8 · p. 10
  • LPT trades time for quality and, at five recursions, costs 4.3–4.8× the decoding FLOPs of a standard Transformer.It spends the same number of recursions on every token; choosing that number per token is future work.Source: §8 · p. 10
  • The evaluation covers models up to 1B parameters on one corpus with 4,096-token contexts, not long-context tasks or very advanced tasks.Analyses use S and M models, and only the detach ablation examines training dynamics.Source: §8 · p. 10

A ARCHITECTURE AND TRAINING DETAILS

This appendix details the model architectures, training setup, and statistical checks. The seed analysis tests whether the reported validation-perplexity gaps persist across runs.

A ARCHITECTURE AND TRAINING DETAILS

  • The decoder-only backbone uses RMSNorm pre-normalization, SwiGLU networks, tied input/output embeddings, and rotary embeddings in full-attention layers.LT and LPT use Ouro’s normalization and recurrent transition; LT2 and H-LPT use the Full+GDN pattern, while PLT adds the input embedding to a shifted normalized prior-recursion output.Source: A ARCHITECTURE AND TRAINING DETAILS · p. 15, Table 2 · p. 15
  • Training uses 4,096-token packed documents, masks attention across document boundaries, and applies either adaptive loss across recursions or final-recursion loss.The implementation concatenates local and shared keys/values at full-attention layers and applies a block mask; all models use the same untuned AdamW schedule and batch setup.Source: A ARCHITECTURE AND TRAINING DETAILS · p. 15

B SEED VARIANCE AND STATISTICAL TESTS

C NORMALIZATION BETWEEN RECURSIONS

The appendix tests whether Ouro’s normalization explains LT’s depth instability or LPT’s advantage. Across the tested alternatives, LPT still has lower perplexity than the best LT normalization tried.

C NORMALIZATION BETWEEN RECURSIONS

  • The tested alternatives remove or retain sandwich norms and the recurrent norm in different ways, while LPT keeps Ouro’s normalization and final loss.These alternatives include pre-norm only, sandwich only, input norm, update norm, and KEEL.Source: C NORMALIZATION BETWEEN RECURSIONS · p. 16
  • LPT’s Ouro-normalized models achieve lower perplexity than the best tested LT normalization by 0.31–0.37 at S and 0.06–0.43 at M.The smallest gap, 0.06 at M with two recursions, is within the seed variation measured at S.Source: C NORMALIZATION BETWEEN RECURSIONS · pp. 16–17, Table 5 · p. 17
  • Alternative normalizations let every finished LT run train with final loss at every size and depth, whereas Ouro-normalized LT degrades or diverges.Update norm wins three settings and KEEL wins the other three; choosing the best tested normalization lowers LT perplexity by 0.13–0.76 relative to the main LT.Source: C NORMALIZATION BETWEEN RECURSIONS · p. 16
  • At S, each alternative LT continues improving through five recursions; at M, even the best LT stops improving after three.Source: C NORMALIZATION BETWEEN RECURSIONS · p. 17
  • LT’s best normalization was selected from six using the same validation perplexity reported, while LPT retained Ouro’s normalization.The authors note that searching LPT’s normalization too could only preserve or lower its perplexity, and keep Ouro’s for the main LT because normalization is not the focus.Source: C NORMALIZATION BETWEEN RECURSIONS · p. 17

D ADDITIONAL ANALYSIS DETAILS

The appendix describes how attention representations and gradient paths are measured, then tests whether shared-cache feedback persists across settings and contributes to learning. The results show persistent shared-memory contribution and worse perplexity when direct shared-cache gradients are detached.

D ADDITIONAL ANALYSIS DETAILS

  • For the attention-representation analysis, each model and layer gets its own PCA basis fitted to pooled, mean-centered query and key vectors before rotary encoding.The panels show the first two components, with centroids connected across recursions; axes are not aligned across models or layers.Source: D ADDITIONAL ANALYSIS DETAILS · p. 17
  • The additional figures cover full-cache key comparisons, query/key geometry by model size and depth, and attention to shared memory across windows and objectives.The shared-memory attention mean averages all layers, although only six matched-relative-depth layers are displayed.Source: D ADDITIONAL ANALYSIS DETAILS · p. 17, Figure 9 · p. 18, Figure 10 · p. 18
  • The gradient probe differentiates final-recursion cross-entropy at matched positions and sequences, then separates shared-cache and hidden-state contributions to the first recursion’s gradient.The probe uses 32 validation sequences and excludes embeddings, readout, recurrent normalization, and the halting gate from measured parameters.Source: D ADDITIONAL ANALYSIS DETAILS · pp. 17–19
  • The shared-cache contribution remains strong with a full prediction window and adaptive loss, consistently across S and M.This indicates the compact comparison does not depend on final loss or a short LPT prediction window.Source: D ADDITIONAL ANALYSIS DETAILS · p. 19, Figure 12 · p. 19
  • The shared-KV-to-hidden-path gradient-norm ratio increases with recursion depth at initialization and is larger after training.Because deeper models add shared-cache reads, the depth trend alone does not establish that feedback helps learning; the detachment experiment tests that separately.Source: D ADDITIONAL ANALYSIS DETAILS · p. 20, Figure 13 · p. 20
  • Detaching shared KV reads raises perplexity in every comparison, supporting a role for shared-cache gradients in learning and stability.The experiment does not distinguish shorter gradient paths from changes in gradient scale or direction.Source: D ADDITIONAL ANALYSIS DETAILS · p. 20, Table 6 · p. 20

E EXPERIMENTS AT NINE RECURSIONS

At nine recursions, H-LPT retains an advantage over LT2, while LPT still exceeds LT across sizes; GPU decoding measurements show that memory savings support larger batches, with throughput tradeoffs.

E EXPERIMENTS AT NINE RECURSIONS

F DECODING SPEED AND MEMORY ON A GPU

Loading 2610.02383v1…