BeNice! Closing Looping's Inference-Memory Gap via KV-Sharing

Hayden Prairie; Jacob M. Ryan; Taylor Berg-Kirkpatrick; Daniel Y. Fu

cs.LGOct 22, 2026

BeNice! Closing Looping's Inference-Memory Gap via KV-Sharing

Authors: Hayden Prairie, Jacob M. Ryan, Taylor Berg-Kirkpatrick, Daniel Y. Fu

Organizations: University of California, San Diego · Together AI

Abstract

Looped architectures scale quality while keeping parameter count fixed by repeatedly passing activations through the same block of layers. During inference, each recurrence generates independent KV-caches, causing memory to grow linearly. To address this challenge, prior work has proposed KV-sharing (i.e., sharing KV-caches across depth), finding that, in some cases, the KV-cache of the final recurrence can be shared with the earlier recurrences without quality loss, a property we denote as KV-shareability. However, the mechanisms underlying KV-shareability remain poorly understood, with even identically trained looped architectures not consistently exhibiting it. In this paper, we study how KV-shareability can be enabled in looped models. We propose a two-stage method: in the first stage, we use stochastic depth during pretraining, which induces contractive residual dynamics that drive representations toward a common limit across recurrences. In the second stage, we propose BeNice!, a post-training method that broadcasts the final recurrence's KV-cache into the first recurrence, training it to produce compatible residuals while preserving the original model's quality. As a result, BeNice! enables KV-shareability while matching standard finetuning's unshared validation perplexity. Using BeNice!, we show that KV-sharing is complementary to alternative KV-compression methods, composing with memory-efficient sequence mixers (e.g., multi-latent attention, Kimi delta attention) and low-precision KV-cache quantization (e.g., fp8, mxfp8, nvfp4). Finally, in simulated low-memory, long-context inference, KV-sharing improves cache hit rates by up to 31.4%, increases throughput by up to 23.6%, and reduces p90 time-to-first-token by up to 83.6%.

Figures & tables

Figure 1: Stochastic Depth, BeNice! , and KV-Shareability. ( Left ) During pretraining, stochastic depth induces contractive residual dynamics. ( Middle ) BeNice! post-trains by injecting the final recurrence KV-caches into the first recurrence. Together, they enable KV-shareability in looped models, retaining only the final recurrence’s KV-cache without degrading quality ( Right ).

	Model	Lambada BPB ( $\downarrow$ )			Core ( $\uparrow$ )			Core-Extended ( $\uparrow$ )
	Model	Unshared	Shared	$\mathbf{\Delta}$	Unshared	Shared	$\mathbf{\Delta}$	Unshared	Shared	$\mathbf{\Delta}$
Fixed-Point	Parcae 140M	0.962	0.962	0.000	14.07	14.17	0.10	10.13	10.09	-0.04
	Parcae 370M	0.762	0.764	0.002	20.47	19.85	-0.62	13.23	13.24	0.01
	Parcae 770M	0.645	1.460	0.814	25.05	21.10	-3.94	15.69	13.95	-1.74
	Parcae 1.3B	0.601	2.281	1.680	28.76	21.37	-7.40	16.73	13.48	-3.25
	Huginn 3.5B	0.429	0.429	0.000	40.97	40.53	-0.44	29.06	28.61	-0.44

Table 1: Emergence of KV-shareability in Looped Models. Quality of looped architectures with and without a shared KV-cache. Observe that KV-shareability is not universal to all looped architectures ( Bolded models have KV-shareability ). See Appendix C for evaluation details.

Figure 2: Residuals and Attention Convergence. Relative step-to-step change of the residual stream (

\lVert h_{i}-h_{i-1}\rVert/\lVert h_{i-1}\rVert

, left ) and attention pattern (

\lVert A_{i}-A_{i-1}\rVert/\lVert A_{i-1}\rVert

, right ) across recurrence depth

T

. Stochastic-depth training exhibits fixed-point convergence, while static-depth does not.

Figure 4: KV-Sharing Failures Converge to Incorrect Fixed Points. We visualize fixed points of several tokens after $T=8$ test-time recurrence with unshared and shared KV-caches. Models with KV-shareability (e.g., 140m and 370m) converge to the same fixed point under both setups, while models without KV-shareability (e.g., 770M and 1.3B) do not. See Appendix H for extended details.

	Model	Core ( $\uparrow$ )			Core-Extended ( $\uparrow$ )
	Model	Unshared	Shared	$\mathbf{\Delta}$	Unshared	Shared	$\mathbf{\Delta}$
140M	Parcae	15.50	14.97	-0.53	10.43	10.10	-0.33
140M	+ BeNice!	15.47	15.61	+0.14	10.28	10.19	-0.09
370M	Parcae	20.74	15.16	-5.58	12.99	9.68	-3.31
370M	+ BeNice!	23.90	23.51	-0.39	14.39	14.24	-0.15
770M	Parcae	30.45	26.92	-3.52	19.22	16.67	-2.55

Table 4: BeNice! Maintains Quality and Enables KV-Shareability. We train Parcae models [ 50 ] and then post-tune with BeNice! . BeNice! maintains and improves base model quality, while fully enabling KV-sharing. Best results are bolded .

	Model	Val. PPL ( $\downarrow$ )			Core ( $\uparrow$ )			Core-Extended ( $\uparrow$ )
	Model	Unshared	Shared	$\mathbf{\Delta}$	Unshared	Shared	$\mathbf{\Delta}$	Unshared	Shared	$\mathbf{\Delta}$
770M	Standard FT	11.86	21.61	+9.75	24.74	23.42	-1.32	15.34	14.54	-0.80
770M	BeNice!	11.85	11.86	+0.01	25.15	25.19	+0.04	15.35	15.29	-0.06
1.3B	Standard FT	10.79	18.52	+7.73	29.01	27.49	-1.51	17.23	16.66	-0.58
1.3B	BeNice!	10.79	10.78	-0.01	28.63	28.67	+0.04	17.32	17.41	+0.09

Table 5: Comparing BeNice! with Standard Finetuning. We finetune public checkpoints of Parcae [ 50 ] on three different seeds, evaluating on fineweb-edu [ 49 ] validation and Core and Core Extended [ 38 ] . BeNice! matches the quality of standard finetuning, while enabling KV-shareability.

Figure 6: KV-Shareability with Memory-Efficient Sequence Mixers. We train Parcae [ 50 ] models with MHA [ 68 ] , MLA [ 13 ] , and KDA [ 65 ] . We find that BeNice! allows memory-efficient sequence mixers to have KV-shareability, reducing inference overhead.

Figure 7: Quality of KV-Sharing and -Quantization. We report changes in climbmix PPL. [ 16 ] , Core, and Core-Extended [ 38 ] after quantizing the KV-shared cache across MHA [ 68 ] model scales. KV-sharing maintains quality across cache precisions.

Figure 10: Throughput Frontier of KV-Sharing. ( Left ) Under increasing offered load, KV-Sharing improves throughput and reduces TTFT over unshared looping. ( Right ) Under high concurrency load, KV-sharing improves overall throughput at matched decode rates.

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Notation	Description
$d_{h}$	Hidden dimension of the model
$t$	Discrete temporal state axis of $\mathcal{R}$ on $\mathbb{N}$
$\gamma$	Hyperparameter used to select the strength of the auxiliary loss in BeNice!
$\mathcal{P}$	Initial prelude block of a recurrent architecture
$\mathcal{R}$	Middle recurrent block of a recurrent architecture
$\mathcal{C}$	Final coda block of a recurrent architecture

Appendix

Table 6: Glossary of notation and terminology. ( Top ) Frequently used dimensions for tensors. ( Middle ) Definition of Parcae blocks, tensors, and distributions used to express recurrent depth models. ( Bottom ) KV-cache compression definitions.

Category	Task	Type	Shots	KV Protocol	CORE
Understanding	HellaSwag (0-shot) [ 74 ]	MC	0	QA	✓
	HellaSwag (10-shot) [ 74 ]	MC	10	QA	✓
	LAMBADA [ 47 ]	LM	0	QA/PPL	✓
	Winograd WSC [ 6 ]	S	0	QA	✓
	WinoGrande [ 55 ]	S	0	QA	✓
	BIG-Bench Language ID [ 60 ]	MC	10	QA	✓

Appendix

Table 7: Downstream evaluation tasks. Tasks marked with ✓ belong to the 22-task CORE suite; all 46 tasks belong to CORE-Extended. Under the QA KV-sharing protocol, the context is prefilled using the final-recurrence cache and candidate continuations are scored against that cache. Under the PPL protocol, used for the separate LAMBADA BPB evaluation, the complete sequence is streamed token-by-token so that every prediction reads only final-recurrence KVs. Type denotes the task scoring rule: MC uses mean candidate NLL, S uses schema-based NLL, and LM requires an exact greedy continuation match.

Figure 11: KV-sharing Error Subspace is Low Rank. The variance within KV-sharing error subspace is well explained with a few directions.

Figure 12: Fixed-point convergence of pretrained looped models. Relative changes in the recurrent residual state and attention pattern decrease with recurrence depth for the released Parcae models and Huginn.

Figure 13: Training Curves for Stochastic- and Static-Depth Models. Training loss as a function of training tokens across model scales. The curves closely overlap, showing that stochastic- and static-depth training exhibit similar optimization dynamics.

	Training	Val. PPL ( $\downarrow$ )			CORE ( $\uparrow$ )
	Training	Unshared	Shared	$\mathbf{\Delta}$	Unshared	Shared	$\mathbf{\Delta}$
140m	Stochastic	25.68	25.67	-0.02	6.41	6.48	+0.08
140m	Static	24.48	25.79	+1.31	9.17	8.55	-0.62
370m	Stochastic	17.34	17.36	+0.02	14.49	14.36	-0.13
370m	Static	17.30	20.00	+2.70	15.00	14.22	-0.78
770m	Stochastic	14.03	14.04	+0.01	18.91	18.78	-0.13

Appendix

Table 8: Quality under stochastic- and static-depth training. We compare the standard per-recurrence cache (Unshared) with final-recurrence KV sharing (Shared).

\Delta

is Shared minus Unshared.

Model	$\mathbf{\mathbb{E}[\\|h_{8}^{(\text{shared})}-h_{8}^{(\text{unshared})}\\|_{2}]}$	$\mathbf{\mathbb{E}[\\|h_{8}^{(\text{unshared})}\\|_{2}]}$	$\mathbf{\mathbb{E}\left[\frac{\\|h_{8}^{(\text{shared})}-h_{8}^{(\text{unshared})}\\|_{2}}{\\|h_{8}^{(\text{unshared})}\\|_{2}}\right]}$
140M	0.4588	56.3	0.82
370M	2.2766	135.2	1.69
770M	186.35	153.4	134.8
1.3B	342.65	301.6	177.5

Appendix

Table 9: Mean difference between shared and unshared final latents, mean unshared latent norm, and mean relative difference.

Model	Baseline: Mean Relative Difference	BeNice! : Mean Relative Difference
140M	0.0526	0.0059
370M	0.1304	0.0059
770M	0.0957	0.0070
1.3B	0.1253	0.0049

Appendix

Table 10: Mean Relative Difference of Fixed Points With and Without KV-sharing Before and After Post-training with BeNice! . Mean relative difference at each model scale for the baseline and BeNice! .

	Model	Val. PPL ( $\downarrow$ )			Core ( $\uparrow$ )			Core-Extended ( $\uparrow$ )
	Model	Unshared	Shared	$\mathbf{\Delta}$	Unshared	Shared	$\mathbf{\Delta}$	Unshared	Shared	$\mathbf{\Delta}$
770M	Parcae	11.87	34.20	+22.33	25.05	21.10	-3.94	15.69	13.95	-1.74
	+ Standard FT	11.86	21.61	+9.75	24.74	23.42	-1.32	15.34	14.54	-0.80
	+ BeNice!	11.85	11.86	+0.01	25.15	25.19	+0.04	15.35	15.29	-0.06
	+ BeNice! †	11.89	11.89	+0.00	24.69	24.71	+0.02	15.32	15.19	-0.13
1.3B	Parcae	10.78	1.6e4	+1.6e4	28.78	21.33	-7.45	16.73	13.48	-3.25

Appendix

Table 11: Quality of BeNice! Finetuning Extended. We provide the extended results of finetuning using BeNice! and

\textsc{BeNice!}^{\dagger}

. Both enable KV-sharing while maintaining quality.

	Model	Val. PPL ( $\downarrow$ )
	Model	Unshared	Layer $>1$ , Recurrence $>1$	Naive Patching	Shared
MHA	140M	15.51	15.50	15.52	16.99
	370M	11.57	11.56	11.56	3723
	770M	9.78	9.80	10.49	36.66
	1.3B	8.58	8.92	12.77	9.42
MLA	140M	15.56	15.54	15.55	15.55

Appendix

Table 12: KV-Sharing Failures on New Models. Preserving first-layer KVs from the first recurrence matches the unshared baseline across model sizes and sequence mixers. Mean patching fails to fully recover quality for MLA at 770M. KV-sharing fails for MLA models as size increases. Surprisingly, KDA models have KV-sharing out-of-the-box.

Figure 14: Transferability of Learning Rate. Optimal learning rate transfers strongly across deeper recurrences for all model backbones.

Figure 15: Transferability of Batch Size. Optimal batch size transfers strongly across deeper recurrences for all model backbones.

Hyperparameter	Model	2 $\iff$ 4	2 $\iff$ 8	4 $\iff$ 8	Mean	$\Delta$ BPB ( $\times 10^{-3}$ )
Learning Rate	140M	0.950	0.967	0.900	0.939	41.805
Learning Rate	370M	0.943	1.000	0.943	0.962	26.483
Learning Rate	770M	0.829	0.829	1.000	0.886	18.373
Muon Momentum	140M	0.900	0.700	0.900	0.833	37.269
Muon Momentum	370M	1.000	1.000	1.000	1.000	25.083
Warmup Fraction	140M	0.600	0.900	0.700	0.733	22.614

Appendix

Table 13: Spearman Rank Correlation of Hyperparameter curves across Recurrence. Here

x\iff y

is the Spearman correlation between the validation BPB hyperparameter curves of a model trained with

\mu_{\text{rec}}=x

and a model trained with

\mu_{\text{rec}}=y

. Furthermore,

\Delta

BPB indicates the validation BPB span across all (recurrence schedule

\times

hyperparameter values) points for that specific parameter class. While Spearman correlation is not

1

in all transfers, we observe that low Spearman correlation coincides with a lower

\Delta

BPB span, which resolves to slight noise on a relatively flat landscape for hyperparameter tuning.

Figure 16: Transferability of Muon Momentum. We observe optimal muon momentum to transfer well across recurrences, with some minor variation in relatively flat regions.

Figure 17: Transferability of Warmup Ratio. We observe optimal warmup ratios to transfer well across recurrences, with some minor variation in relatively flat regions.

		Small (140M)	Medium (370M)	Large (770M)	XLarge (1.3B)
Architecture	Layers in $\mathcal{P}$	2	4	6	8
	Layers in $\mathcal{R}$	2	4	6	8
	Layers in $\mathcal{C}$	2	4	6	8
	$d_{\text{model}}$	768	1,024	1,280	1,536
	$d_{\text{intermediate}}$	3,072	4,096	5,120	6,144
MHA	Attention	Causal Self-Attention [ 68 ] w/ QK-Norm [ 26 ]

Appendix

Table 14: Model definitions of baseline Parcae [ 50 ] , Parcae with MLA, and Parcae with KDA.

Optimizer	Parameters
AdamW [ 36 ]	Token embeddings ( wte )
	LM head ( lm_head )
	Normalization layers ( RMSNorm )
	Value embedding gates ( ve_gate )
	All 1D parameters
	Injection parameters ( $\bm{A}$ , $\Delta$ , $\bm{B}$ )

Appendix

Table 15: Optimizer parameter group assignment for Parcae.

	Small (140M)	Medium (370M)	Large (770M)	XLarge (1.3B)
AdamW and Muon Learning Rate	$2e^{-2}$	$1.2e^{-2}$	$6e^{-3}$	$5e^{-3}$
Warmup Ratio	0.05
Constant Decay Ratio	0.2
Gradient Clipping	1.0
Muon Momentum	0.95
Muon Weight Decay	0.3

Appendix

Table 16: Pretraining Hyperparameters for climbmix trained Parcae models (including MLA and KDA variants).

KV-Cache Reuse

Key-value (KV) cache reuse lets language models reuse previously computed attention states to cut repeated prefill work and serving latency.

See all 73 papers

Sep 24, 2026

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

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.

Wanqi Yang, Shiwei Liu

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

Efficient Transformer Inference Recurrent Transformers

Ask PDF

Oct 1, 2026

The Surprising Effectiveness of Shared Memory in Looped Transformers

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.

Giovanni Monea, Keshav Ramji, Yousef El-Kurdi +4

Cornell University · IBM Research

Memory-Efficient Inference Recurrent Transformers

Ask PDF

Oct 7, 2026

ResidualQuant: KV Cache Quantization for Looped Transformers with 2-Bit Residuals

Looped Transformers improve parameter efficiency by repeatedly applying shared Transformer blocks over multiple recurrent loops, increasing computational depth without increasing the parameter count. However, KV cache memory still scales with the number of loops, becoming a key memory bottleneck that limits batch size and inference throughput. KV cache quantization can alleviate this bottleneck, but existing methods often suffer substantial accuracy degradation at aggressive low-precision regimes. We observe that looped Transformers offer a unique opportunity: KV states across loops are highly similar. Based on this observation, we propose ResidualQuant, which uses the final-loop KV states as a reference and represents the remaining loops with low-precision residuals. Our method further combines least-square scaling and rotations applied to the residuals, as well as loop-wise mixed precision, to enable accurate quantization down to INT2 while retaining efficient reconstruction. Across multiple looped Transformer models and mathematical reasoning and code generation benchmarks, ResidualQuant consistently improves the accuracy-memory tradeoff over state-of-the-art rotation-based KV quantization. In particular, our method retains accuracy close to BF16 under mixed-precision settings while reducing theoretical KV storage by 80.7%, achieving up to 13.0% higher accuracy than the rotation-based baseline at the same memory budget. On an RTX 5090, the reduced KV memory traffic improves fixed-batch decode throughput by up to 2.73x, while the smaller memory footprint enables up to 2x larger batches, improving peak throughput by up to 4.15x.

Heejun Kim, Junyoung Lee, SangLyul Cho +3

KAIST · Yonsei University · Seoul National University

KV-Cache Quantization Efficient Transformer Inference Low-Bit Quantization Recurrent Transformers

Ask PDF

Abstract

KV-Cache Reuse

Explore similar work

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

The Surprising Effectiveness of Shared Memory in Looped Transformers

ResidualQuant: KV Cache Quantization for Looped Transformers with 2-Bit Residuals