LSTM-Style Looped Transformers with Shared KV Caches

Andreas Kirsch

cs.LGOct 23, 2026

LSTM-Style Looped Transformers with Shared KV Caches

Authors: Andreas Kirsch

Abstract

Recurrent transformers reuse weights to perform deep computation with fewer parameters. We study an LSTM-style transformer that also shares one key--value cache across its recurrent core. At width 1024, it uses 40.4% fewer parameters than an untied control at nearly matched forward FLOPs, and reduces persistent cache sets from fifteen to four. Parallel training supplies cache entries from a standard-attention pass, whereas continuation reads entries produced recursively using shared attention. We address this mismatch by aligning the two training passes' cache entries with a consistency loss during the final 10% of a fixed training-token budget. Across three model widths on FineWeb, this removes 81--89% of the continuation-loss penalty relative to full-cache evaluation. The improvement replicates on disjoint documents for two training seeds. Across four widths and 1.05--4.19B training tokens, the loss cost of weight tying is smallest at lower token budgets, suggesting an opportunity to use recurrent computation when data and model storage are limited.

Figures & tables

Panel of Figure 1 — Figure 1: Parameter and cache sharing have distinct costs. (a) The four core blocks have distinct weights, reused across iterations; preceding tokens supply one common donor cache. Prologue and epilogue retain their own caches. (b) Widths 512, 768, and 1024 after 4.19B training tokens, evaluated on the same 1,280 document continuations with 2,048 prompt tokens and 512 scored targets. Repaired models add donor consistency during the final training stage. The FLOPs panel estimates one inference forward, not latency or the two-pass training cost (Appendix D ). All points use one training seed.

Width	Model	Params (M)	Forward FLOPs (M/token)	Sets	KV (KiB/token)	Full-cache loss $\downarrow$	Shared-cache loss $\downarrow$
512	Untied	70.01	178.0	15	15	3.3679	—
	Two-pass	46.93	181.1	4	4	3.4395	3.5589
	Repaired	46.93	181.1	4	4	3.4420	3.4570
768	Untied	135.98	328.9	15	22.5	3.2393	—
	Two-pass	85.24	336.0	4	6	3.2827	3.3347
	Repaired	85.24	336.0	4	6	3.2839	3.2938

Table 1: Continuation loss, model size, forward computation, and persistent KV storage after 4.19B training tokens. Full-cache and shared-cache losses score identical targets under ordinary prompt prefill (Section 3 ). Repaired models repeat their parents’ final training stage with donor consistency, retaining the same token budget. FLOPs use the reference attention lengths in Appendix D ; two-pass training costs about twice as much at fixed width and tokens. Sets and KV storage describe each model’s serving layout: shared-cache columns use four sets; recurrent full-cache evaluation requires fifteen. BF16 KV storage is calculated from tensor dimensions for equal-length cache sets, excluding scratch space; window-sized allocations are discussed in Section 4.3.2 . Lower loss is better.

Figure 2: Loss across widths and training-token budgets. The three models are separately trained. Standard and parallel shared-cache validation score the same 1,280 chunks (10,485,760 targets). Each width has a 15-layer untied control approximately matched to one recurrent forward; two-pass training costs more. Points use seed 0; Appendix F repeats the width-768 two-pass/untied comparison over three seeds. Appendix E replots the losses against reference forward FLOPs. These parallel losses do not measure shared-cache continuation.

Update	Adapter	Params (M)	Loss $\downarrow$	$\Delta$ loss $\downarrow$
Gated carry	No	46.61	3.4517	—
Plain overwrite	No	45.43	3.4838	$+0.0321$ (0.0003)
Gated carry	Yes	47.79	3.4518	$+0.0001$ (0.0003)
Plain overwrite	Yes	46.61	3.4683	$+0.0166$ (0.0003)
Huginn-style	Yes	46.61	3.4745	$+0.0228$ (0.0003)

Table 2: Recurrent update and input-adapter controls. P1E2 at width 768, three iterations and 1.05B training tokens, one training seed. Losses use standard attention on 10,485,760 targets. Parentheses are paired block-bootstrap SEs of the difference from the cell, conditional on these checkpoints. Parameter counts exclude embeddings.

Figure 3: Excess loss during shared-cache continuation. Width 768, 4.19B training tokens, seed 0, before and after late donor consistency. Each curve subtracts the same checkpoint’s full-cache loss at the same positions. Both settings use the original 1,280 documents; prompt and continuation lengths differ together. The repair reduces the excess loss throughout the continuation, without making the two cache computations equivalent.

Width	$\Delta$ continuation loss $\downarrow$	$\Delta$ full-cache loss $\downarrow$	Residual penalty
512	$-0.1019$ (0.0022)	$+0.0025$ (0.0002)	$+0.0150$
768	$-0.0409$ (0.0010)	$+0.0012$ (0.0003)	$+0.0099$
1024	$-0.0565$ (0.0013)	$+0.0011$ (0.0003)	$+0.0075$

Table 3: Effect of late donor consistency. Repaired minus untreated checkpoint at 4.19B tokens, seed 0, on the primary 2,048-prompt/512-target setting. Parentheses are paired block-bootstrap SEs. The residual penalty subtracts the repaired model’s own full-cache loss from its shared-cache continuation loss. Absolute losses and resources are in Table 1 .

Figure 4: Fixed depth outperforms the tested stopping rule. The width-768 model is trained over depths 3–6 on 2.10B tokens, with and without late donor consistency. Losses use shared-cache continuation on the primary 1,280 spans. Fixed depths are three, four and six; per-token stopping uses successive-state cosine

\tau=0.98

, minimum depth two and cap six. The horizontal axis counts selected passes, not executed GPU computation. There is one training seed.

Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Endpoint updates	Decay starts	Training tokens
4,000	3,600	1,048,576,000
8,000	7,200	2,097,152,000
16,000	14,400	4,194,304,000

Appendix

Table 4: Training endpoints. Token counts include the complete path to each endpoint; extra branches consume additional experimental computation but do not add to that path.

Width	Trained model	Total params	Non-emb.	$\Delta P$ (%)	Forward FLOPs	$\Delta F$ (%)	Training FLOPs
384	Untied	43.66	24.34	—	115.8	—	347.4
384	Standard cell	30.97	11.66	-29.1	117.6	+1.5	352.7
384	Two-pass cell	30.97	11.66	-29.1	117.6	+1.5	687.5
512	Untied	70.01	44.25	—	178.0	—	534.0
512	Standard cell	46.93	21.18	-33.0	181.1	+1.8	543.4
512	Two-pass cell	46.93	21.18	-33.0	181.1	+1.8	1058.2

Appendix

Table 5: Resources at three recurrent iterations. Parameters are millions; FLOPs are millions per token. Forward is a single standard forward for every row; training counts both passes for the two-pass objective. The repair shares the two-pass row’s base cost.

\Delta P

and

\Delta F

are percentage differences in total parameters and forward FLOPs relative to the same-width untied control.

Width	Params (M) T / U	Training FLOPs T / U	Compute ratio T/U	T loss $\downarrow$	U loss $\downarrow$	$\Delta$ loss T $-$ U
384	30.97 / 43.66	1.442 / 1.457	0.989	3.6257	3.4819	$+0.1439\ (0.0006)$
512	46.93 / 70.01	2.219 / 2.240	0.991	3.4967	3.3710	$+0.1257\ (0.0006)$
768	85.24 / 135.98	4.112 / 4.139	0.993	3.3448	3.2394	$+0.1054\ (0.0008)$
1024	136.20 / 228.49	6.670 / 6.706	0.995	3.2514	3.1557	$+0.0957\ (0.0008)$

Appendix

Table 6: Approximately matched estimated training computation, using existing seed-0 endpoints. T denotes the two-pass model at 2.097B tokens and U the untied model at 4.194B. Parameters include embeddings; computation is total training FLOPs in units of

10^{18}

. Losses score the same 1,280 chunks (10,485,760 targets): T uses parallel shared-cache validation before repair, U standard validation. Parentheses are paired block-bootstrap SEs, conditional on the checkpoints.

Figure 5: Loss versus reference forward computation across training-token budgets. The same checkpoints and losses as Figure 2 , with widths 384, 512, 768, and 1024 in increasing compute order. The horizontal axis counts one standard forward under Appendix D ’s attention-length convention, as an inference compute reference; it does not count two-pass training or measure serving latency. Untied and single-pass cell losses use standard validation; two-pass cell losses use parallel shared-cache validation before consistency repair. All points use training seed 0 and the same 10,485,760 targets.

Width	Tokens	Untied	Std. cell	Two-pass std.	Two-pass shared	$\Delta_{\rm cell}$	$\Delta_{\rm shared}$
384	1.049	3.6880	3.7158	3.7338	3.7329	+0.0278	+0.0449
384	2.097	3.5641	3.6083	3.6271	3.6257	+0.0442	+0.0617
384	4.194	3.4819	3.5408	3.5594	3.5582	+0.0589	+0.0763
512	1.049	3.5759	3.5941	3.6076	3.6035	+0.0182	+0.0276
512	2.097	3.4549	3.4864	3.5011	3.4967	+0.0315	+0.0417
512	4.194	3.3710	3.4138	3.4292	3.4256	+0.0428	+0.0546

Appendix

Table 7: Complete chunk-validation ladder; loss

\downarrow

. Tokens are billions.

\Delta_{\rm cell}

is standard-cell minus untied;

\Delta_{\rm shared}

is two-pass shared minus untied. Each row scores the same 10,485,760 targets.

Tokens	Seed	Untied	Two-pass std.	Two-pass shared	$\Delta$
1.049	0	3.4462	3.4624	3.4571	+0.0109
1.049	1	3.4476	3.4637	3.4587	+0.0111
1.049	2	3.4444	3.4620	3.4569	+0.0125
2.097	0	3.3254	3.3504	3.3448	+0.0194
2.097	1	3.3264	3.3531	3.3479	+0.0215
2.097	2	3.3235	3.3512	3.3455	+0.0220

Appendix

Table 8: Three training seeds at width 768, on the common chunk evaluation. Loss

\downarrow

;

\Delta

is two-pass shared minus untied. Tokens are billions.

Seed	Parent shared	Repaired shared	Shared difference	Full-cache difference	Untied
0	3.3346	3.2911	$-0.0434\ (0.0011)$	$+0.0010\ (0.0002)$	3.2390
1	3.3235	3.2936	$-0.0298\ (0.0006)$	$+0.0014\ (0.0003)$	3.2336

Appendix

Table 9: Held-out width-768 confirmation at 4.19B tokens, with 2,048 prompt tokens and 512 targets. Differences compare the repair to its own untreated parent; parentheses are paired bootstrap SEs. All losses are in nats per target (

\downarrow

Auxiliary objective	Full	Shared	Penalty	$\Delta$ chunk std.
None	3.2827	3.3353	+0.0526	+0.0000
Donor + epilogue, $\gamma=0.1$	3.2826	3.3029	+0.0203	+0.0001
Donor + epilogue, $\gamma=0.3$	3.2834	3.2956	+0.0122	+0.0010
Donor only, $\gamma=0.1$	3.2832	3.3001	+0.0169	+0.0004
Donor only, $\gamma=0.3$	3.2839	3.2938	+0.0099	+0.0013
Donor only, $\gamma=1.0$	3.2884	3.2902	+0.0018	+0.0052

Appendix

Table 10: Width-768 auxiliary-loss ablations on the same 1,600-update decay branch. Full and shared losses score the primary continuation targets (

\downarrow

). Penalty is shared minus full within that model. The last column is the change in chunk-standard loss relative to the no-auxiliary branch control; its preset quality limit is

+0.005

Update	$\gamma$	Params (M)	Fwd MFLOPs/tok	Cold loss $\downarrow$	$\Delta$ loss $\downarrow$ (SE)
Cell	0	46.933	181.1	3.7961	—
Cell	0.3	46.933	181.1	3.6587	$-0.1374\;(0.0026)$
Plain	0	46.408	178.0	3.7658	—
Plain	0.3	46.408	178.0	3.6844	$-0.0814\;(0.0013)$
Injection	0	46.932	181.1	3.6945	—
Injection	0.3	46.932	181.1	3.6674	$-0.0271\;(0.0005)$

Appendix

Table 11: Recurrent-update repair pilot: width 512, P1E2,

R=3

, 1.048576B training tokens and four persistent cache sets (4 KiB/token). Cold loss scores 512 targets after 2,048 prompt tokens per document.

\Delta

is repair minus the corresponding untreated branch, with paired bootstrap SE. Reference forward cost follows Appendix D .

Checkpoint	$k=1$	$k=2$	$k=3$	$k=4$	$k=5$	$k=6$
Untreated	3.2568	3.2510	3.2994	3.3246	3.3290	3.3313
Late donor consistency	3.2579	3.2515	3.2596	3.2640	3.2663	3.2671

Appendix

Table 12: Recorded donor-iteration diagnostic at width 768 and 4.194B training tokens, seed 0: 64 chunks of 8,192 targets (524,288 total). Each

k

is a new whole-sequence forward with three core iterations, not an additional core iteration within one forward. Loss

\downarrow

; the first forward uses standard attention.

Producer	Training objective	Tokens (B)	Train cost	Fwd MFLOPs/tok	Cold loss
Late	Two-pass + consistency	1.049	$1.000\times$	181.1	3.6587
Early	Two-pass	1.049	$1.001\times$	181.7	3.7362
Early	Shared CE	1.049	$0.515\times$	181.7	3.7247
Early	Shared CE	2.036	$1.000\times$	181.7	3.6240

Appendix

Table 13: Early versus repaired late cache production. All models have 46.933M parameters and four persistent cache sets (4 KiB per token in BF16). Training cost is relative to the repaired late model. Reference forward FLOPs use Appendix D ’s convention plus the early producer’s extra K/V projection; they are not measured latency or prefill cost. Cold loss

\downarrow

Producer	Tokens (B)	Train FLOPs ( $10^{18}$ )	Fwd MFLOPs/tok	Cold loss $\downarrow$
Late, two-pass + consistency	4.194	4.4382	181.144	3.4494
Early, shared CE	2.036	1.1095	181.668	3.5903

Appendix

Table 14: Producer comparison on fresh documents, separate from the pilot selection set. Both width-512 P1E2 checkpoints have 46.933M parameters,

R=3

and four cache sets (4 KiB/token). Cold loss scores 655,360 targets with 2,048-token prompts. Training FLOPs exclude auxiliary-loss overhead; forward FLOPs are reference estimates, not measured serving latency.

Sweeps $n$	$\gamma$	Tokens (B)	Train compute	Handoff	Cold	Chunk std.
2	0	1.049	$1.00\times$	3.7122	3.7961	3.6081
2	0.3	1.049	$1.00\times$	3.6523	3.6587	3.6101
4	0	1.049	$1.58\times$	3.6684	3.6522	3.6100
4	0.3	1.049	$1.58\times$	3.6733	3.6539	3.6113
4	0	0.663	$1.00\times$	3.7592	3.7484	3.7173
4	0.3	0.663	$1.00\times$	3.7632	3.7515	3.7187

Appendix

Table 15: Cache-refinement training ablation at width 512. Training compute is relative to the two-sweep, 1.048576B-token reference (1,109.560 estimated PFLOPs), using the accounting in Appendix D . Four-sweep compute-matched runs use 2,529 updates (a

+0.0183\%

compute mismatch); the other runs use 4,000. All sequential evaluations retain

R=3

and one shared core cache; they do not execute additional sweeps. Loss

\downarrow

Tokens, $\gamma$	$k=1$	$k=2$	$k=3$	$k=4$	$k=5$	$k=6$	Cold
1.049B, 0	3.6001	3.6525	3.6482	3.5882	3.6413	3.6215	3.6073
1.049B, 0.3	3.6016	3.6639	3.6501	3.5886	3.6526	3.6376	3.6099
0.663B, 0	3.7170	3.7546	3.7462	3.7053	3.7328	3.7249	3.7175
0.663B, 0.3	3.7185	3.7595	3.7523	3.7059	3.7475	3.7388	3.7213

Appendix

Table 16: All six donor-refinement readouts for the four-sweep-trained models, on 64 chunks of 8,192 targets (524,288 total). The trained shared objective uses sweep 4. Cold sequential loss scores the same chunks. Loss

\downarrow

; these chunk losses are separate from the continuation losses above.

Arm	Full-cache continuation	Shared continuation	Chunk standard
Parent	3.2827	3.3347	3.2728
Reset control	3.4093	3.4602	3.3917
KL only, $\beta=1$	3.4207	3.4272	3.4035
K/V, $\gamma=0.1$	3.4099	3.4287	3.3929
Both	3.4220	3.4199	3.4045

Appendix

Table 17: Optimizer-reset pilot: primary continuation losses and separate chunk-standard loss (

\downarrow

). The parent has 4.194B training tokens; pilot arms have 4.456B.

Update	Adapter	1.049B tokens	2.097B tokens	4.194B tokens
Gated carry	No	3.4517	3.3427	3.2643
Plain overwrite	No	3.4838	—	—
Gated carry	Yes	3.4518	—	—
Plain overwrite	Yes	3.4683	—	—
Huginn-style	Yes	3.4745	3.3649	3.2891
Untied	—	3.4462	3.3254	3.2394

Appendix

Table 18: Update-rule comparisons across training-token budgets. Width 768, with three recurrent iterations for P1E2 and a 15-layer untied control. Columns give standard validation loss (

\downarrow

) after the final learning-rate decay at each budget, on the same 1,280 chunks (10,485,760 targets). One training seed per arm; Huginn-style losses use the same evaluation-noise draw as Table 2 . A dash in a loss column denotes an experiment not run at that budget. Parentheses in the final row are paired block-bootstrap SEs, conditional on the checkpoints.

Update	Backpropagation	$R=4$	$R=8$	$R=16$	$R=32$	$R=64$
Cell	Full	3.6656	3.6112	3.6028	3.6065	3.6287
Cell	Last eight	3.7076	3.6748	3.6731	3.6822	3.7107
Huginn-style	Full	3.6581	3.6052	3.6036	3.6036	3.6037
Huginn-style	Last eight	3.6472	3.5913	3.5896	3.5896	3.5896

Appendix

Table 19: Forced-depth standard losses (

\downarrow

) after uniform-depth training at width 512. Each entry scores 10,485,760 targets; one training seed per arm.

Figure 6: Depth extrapolation with full and truncated gradients, using the models in Table 19 . Training depths range from 4 to 32; depth 64 is extrapolation.

Evaluation $R$	Fixed $R=3$ standard	Fixed $R=3$ shared	Variable $R$ standard	Variable $R$ shared
3	3.3504	3.3448	3.3693	3.3577
4	3.4142	3.4691	3.3513	3.3452
5	—	—	3.3477	3.3436
6	3.7543	3.9704	3.3483	3.3449
8	—	—	3.3551	3.3515

Appendix

Table 20: Parallel losses at width 768 and 2.097B tokens (

\downarrow

). A dash denotes an unrecorded fixed-depth comparison.

Serving rule	Untreated	Whole-run consistency	Late consistency
Fixed $R=3$	3.4187 (3.00)	3.3899 (3.00)	3.3862 (3.00)
Fixed $R=4$	3.4115 (4.00)	3.3803 (4.00)	3.3761 (4.00)
Fixed $R=6$	3.4149 (6.00)	3.3832 (6.00)	3.3779 (6.00)
Halting, cap 6	3.4151 (5.17)	3.3826 (4.18)	3.3816 (4.75)

Appendix

Table 21: Primary shared-cache continuation losses for the variable-depth model, before repair, with whole-run EMA-normalized consistency, and with late fixed-scale consistency (

\downarrow

). Numbers in parentheses are mean selected recurrent passes.

Width	Layout	$B=1$	$B=16$	$B=64$	$B=160$	Peak	Peak $-$ base
768	Cell, standard	9.1	8.6	14.0	29.1	8259	7897
768	Cell, shared	8.3	9.0	13.9	28.9	2980	2618
768	Untied	7.8	7.7	13.5	28.7	8556	7996
1024	Cell, standard	9.2	8.4	17.6	20.7	11136	10561
1024	Cell, shared	8.3	8.4	17.3	37.2	4097	3522
1024	Untied	8.0	8.2	17.3	20.6	11691	10740

Appendix

Table 22: Isolated decoder measurements. Step times are milliseconds (

\downarrow

);

B

is batch size. Peak and peak-minus-base memory are MiB at

B=160

(

\downarrow

Figure 7: Peak allocated memory by batch size for the isolated decoder measurements. Persistent-KV allocation and total peaks are distinct quantities; model weights and scratch computation contribute to the latter.

Comparison	Rule or scope fixed beforehand	Observed outcome
Untied control	Match reference forward FLOPs within 5% before comparing losses.	Fifteen blocks; residual forward mismatch 1.5–2.4%.
Width/token ladder	Widths 384/512/768/1024, three annealed token endpoints, same data and recipe.	Complete; recurrent gaps generally widen with tokens and shrink with width.
Gradient-path control	Cell/injection crossed with full/last-eight gradients; width 512, uniform depth 4–32.	Large truncation penalty for the cell; small full-gradient difference between updates.
Training-seed replication	Three paired seeds for the width-768 two-pass and untied models, with fixed data order.	Shared-minus-untied gap positive in all nine endpoint/seed pairs.
Repair effect and quality	Shared continuation improvement at least 0.01 nats, with paired uncertainty; chunk-standard degradation at most 0.005.	Selected late donor repair meets both at tested widths and adaptive caps; whole-run adaptive treatment exceeds the quality limit.
Branch control	Before comparing auxiliary terms, reproduce parent parallel losses within 0.005 and continuation penalty within 0.01.	Control meets both tolerances.

Appendix

Table 23: Compact comparison and decision-rule record. Outcome statements apply only to the tested scales and evaluation protocols.

Recurrent Transformers

Recurrent Transformers reuse layers or state across steps to extend context, reasoning depth, and sequence modeling with less parameter growth.

See all 180 papers

The Recurrent Transformer: Greater Effective Depth and Efficient Decoding
Apr 23, 2026Costin-Andrei Oncescu, Depen Morwani, Samy Jelassi +3Autoregressive Language Modeling Self-Attention
Harvard University
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
Latent Recurrent Transformer: Architecture Exploration, Training Strategies, and Scaling Behavior
May 26, 2026Zeyi Huang, Xuehai He, LiLiang Ren +8Language Model Pretraining Transformer Inference
Microsoft · University of Washington · University of Wisconsin-Madison

Apr 23, 2026J/K move · Enter open · S save

The Recurrent Transformer: Greater Effective Depth and Efficient Decoding

Costin-Andrei Oncescu, Depen Morwani, Samy Jelassi, Alexandru Meterez+2

Harvard University

Autoregressive Language Modeling Self-Attention Attention Mechanisms Efficient Transformer Inference Recurrent Transformers Efficient Attention+1

Open paper AskPDF arXiv

Transformers process tokens in parallel but are temporally shallow: at position

t

, each layer attends to key-value pairs computed based on the previous layer, yielding a depth capped by the number of layers. Recurrent models offer unbounded temporal depth but suffer from optimization instability and historically underutilize modern accelerators. We introduce the Recurrent Transformer, a simple architectural change where each layer attends to key-value pairs computed off its own activations, yielding layerwise recurrent memory while preserving standard autoregressive decoding cost. We show that the architecture can emulate both (i) a conventional Transformer and (ii) token-to-token recurrent updates under mild assumptions, while avoiding optimization instability. Naively, prefill/training appears bandwidth-bound with effective arithmetic intensity near

1

because keys and values are revealed sequentially; we give an exact tiling-based algorithm that preserves the mathematical computation while reducing HBM traffic from

Θ(N^2)

Θ(N\log N)

, increasing effective arithmetic intensity to

Θ(N/\log N)

for sequence length

N

. On 150M and 300M parameter C4 pretraining, Recurrent Transformers improve cross-entropy over a parameter-matched Transformer baseline and achieve the improvement with fewer layers (fixed parameters), suggesting that recurrence can trade depth for width, thus reducing KV cache memory footprint and inference latency.

Abstract

Recurrent Transformers

Explore similar work

The Recurrent Transformer: Greater Effective Depth and Efficient Decoding

The Surprising Effectiveness of Shared Memory in Looped Transformers

Latent Recurrent Transformer: Architecture Exploration, Training Strategies, and Scaling Behavior