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

Markus Frey; Behzad Shomali; Ali Hamza Bashir; David Berghaus; Joachim Koehler; Mehdi Ali

cs.CLMar 9, 2026

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

Authors: Markus Frey, Behzad Shomali, Ali Hamza Bashir, David Berghaus, Joachim Koehler, Mehdi Ali

Organizations: Lamarr Institute · Fraunhofer IAIS · University of Bonn

Abstract

Chain-of-thought (CoT) prompting enables reasoning in language models but requires explicit verbalization of intermediate steps. Looped transformers offer an alternative by iteratively refining representations within hidden states. This parameter efficiency comes at a cost, as looped models lack the storage capacity of deeper models which use unique weights per layer. In this work, we investigate transformer models that feature both adaptive per-layer looping, where each transformer block learns to iterate its hidden state via a learned halting mechanism, and gated memory banks, that provide additional learned storage. We find that looping primarily benefits mathematical reasoning, while memory banks help recover performance on commonsense tasks compared to parameter and FLOP matched models. Combining both mechanisms yields a model that outperforms an iso-FLOP baseline, with three times the number of layers, across math benchmarks. Analysis of model internals reveals layer specialization: early layers learn to loop minimally and access memory sparingly, while later layers do both more heavily.

Figures & tables

Panel of Figure 1 — Figure 1: Architecture overview. Left: A standard transformer passes hidden states through $L$ unique blocks. Center: Our loop model allows each block to iterate up to $N$ times, with a learned halting mechanism that produces a weighted combination of intermediate states. Per-step scales $\zeta(s_{n})$ are initialized near zero for training stability. Right: The combined model additionally retrieves from local (per-layer) and global (shared) memory banks, gated by learned input-dependent scalars.

Without Memory
Model	Type	CS Acc $\uparrow$	CS BPB $\downarrow$	Math BPB $\downarrow$
Base	IsoPar baseline	0.477	0.859	2.163
Loop-3	Ours ( $N_{\max}{=}3$ )	0.501	0.813	1.687
Loop-5	Ours ( $N_{\max}{=}5$ )	0.503	0.823	1.737
Loop-7	Ours ( $N_{\max}{=}7$ )	0.498	0.832	1.659
IsoFLOP	36-layer	0.523	0.780	1.801

Table 1: Summary of results , averaged across benchmarks within each group. CS = commonsense; BPB = bits per byte (lower is better, see Appendix A.2 for details). Best result per column within each model group is bolded . Full per-benchmark breakdowns are in Appendix A.4 .

Figure 2: Expected number of loop iterations per layer over training. Left: Each curve represents one layer. Early layers (lighter colors) consistently use fewer iterations than later layers (darker colors). Middle: Expected steps at the end of training. Right: All models show a characteristic transition which occurs at approximately the same cross-entropy value across configurations (see Figure 3 in the Appendix for all configurations).

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Expected loop iterations vs. validation cross-entropy for all configurations. Each point represents one evaluation during training; curves are colored by model configuration. Across all looped models, the expected number of iterations begins to increase rapidly once the cross-entropy drops below approximately $3.27\pm 0.59$ . This phase transition is consistent across Loop-3, Loop-5, and Loop-7 configurations, suggesting it depends on the model’s language competence rather than the maximum number of allowed iterations.

Configuration	Loops	Memory	Parameters	Description
IsoPar	×	×	200M	12 Layers
Loop- $N$	✓	×	200M	Adaptive looping, $N_{\max}\in\{3,5,7\}$
IsoFLOP	×	×	332M	36 Layers
IsoPar-M	×	×	210M	Wider FFN
Mem ( $g_{0}$ )	✓	✓	210M	Loop-3 + Memory
IsoFLOP-M	×	×	480M	36 Layers + Wider FFN

Appendix

Table 2: Model configurations. All models use the same base transformer architecture. Loop parameters include per-step scales and halting router weights. Memory parameters include local/global key-value banks and gating networks. Iso-parameter baselines add extra FFN capacity to match the corresponding model’s parameter count. Iso-FLOP baselines use 36 layers to approximate the forward-pass cost of 3-loop models.

Figure 4: Memory gate activations across layers and training. Left: Local memory gate values show high variance across layers while later layers tend to have higher gate activations, and the spread increases over training. Right: Global memory gate values increase during training but converge to a more uniform profile across layers, with activations rising up to approximately layer 5 and then plateauing.

Commonsense Accuracy $\uparrow$
	Without Memory					With Memory
Bench	Base	L3	L5	L7	L3 IF	M-B IP	M -3	M 0	M 3	M-B IF
ARC-C	0.367	0.375	0.492	0.398	0.430	0.398	0.414	0.359	0.398	0.438
ARC-E	0.609	0.672	0.609	0.586	0.688	0.602	0.648	0.672	0.641	0.680
HellaSwag	0.445	0.469	0.445	0.438	0.508	0.453	0.461	0.453	0.461	0.508
Lambada	0.211	0.266	0.289	0.281	0.211	0.227	0.250	0.281	0.242	0.273

Appendix

Table 3: Full results at the final checkpoint. Best result within each group is bolded. Base = standard transformer; L

N

= Loop-

N

; L3 IF = iso-FLOP for Loop-3; M-B IP = memory iso-parameter baseline; M

g_{0}

= memory model with gate init

g_{0}

; M-B IF = memory iso-FLOP baseline.

Commonsense Accuracy $\uparrow$
Bench	Base	L3	L5	L7	L3 IF	M-B IP	M -3	M 0	M 3	M-B IF
ARC-C	0.320	0.281	0.336	0.336	0.359	0.359	0.312	0.305	0.367	0.344
ARC-E	0.547	0.469	0.477	0.523	0.602	0.531	0.508	0.508	0.555	0.586
HellaSwag	0.430	0.391	0.398	0.422	0.414	0.398	0.383	0.391	0.359	0.414
Lambada	0.164	0.156	0.188	0.164	0.148	0.133	0.219	0.203	0.164	0.172
PIQA	0.555	0.609	0.633	0.617	0.539	0.594	0.641	0.609	0.609	0.578

Appendix

Table 4: Results at the early checkpoint (step 5000).

Commonsense Accuracy $\uparrow$
Bench	Base	L3	L5	L7	L3 IF	M-B IP	M -3	M 0	M 3	M-B IF
ARC-C	0.398	0.414	0.453	0.375	0.438	0.312	0.367	0.352	0.359	0.453
ARC-E	0.547	0.617	0.523	0.562	0.625	0.594	0.625	0.578	0.648	0.602
HellaSwag	0.406	0.375	0.406	0.430	0.461	0.391	0.422	0.438	0.430	0.453
Lambada	0.234	0.328	0.297	0.273	0.289	0.250	0.297	0.273	0.289	0.289
PIQA	0.594	0.625	0.609	0.656	0.664	0.633	0.641	0.641	0.656	0.672

Appendix

Table 5: Results at the mid checkpoint (step 20000).

Transformer

Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.

See all 528 papers

The Surprising Effectiveness of Shared Memory in Looped Transformers
Oct 1, 2026Giovanni Monea, Keshav Ramji, Yousef El-Kurdi +4Memory-Efficient Inference Recurrent Transformers
Cornell University · IBM Research
Chain-of-Thought and Compressed Looped Transformers: A Memory-Budget Separation
May 29, 2026Haozhou ZhangCoT Reasoning Persistent Memory for Language Models
Department of Mathematics and Statistics University of Idaho
Loop the Loopies!
Jul 17, 2026Zitian Gao, Yilong Chen, Yihao Xiao +4Mixture-of-Experts Language Models Recurrent Transformers
IQuest Research

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

The Surprising Effectiveness of Shared Memory in Looped Transformers

Giovanni Monea, Keshav Ramji, Yousef El-Kurdi, Luis A. Lastras+3

Cornell University · IBM Research

Memory-Efficient Inference Recurrent Transformers

Open paper AskPDF arXiv

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.

Abstract

Transformer

Explore similar work

The Surprising Effectiveness of Shared Memory in Looped Transformers

Chain-of-Thought and Compressed Looped Transformers: A Memory-Budget Separation

Loop the Loopies!