Looped language models promise a compelling trade-off: scale model quality by recycling the same parameters across multiple passes without growing the parameter count. The catch is that each pass through the recurrent block generates its own key-value cache during inference. With eight or sixteen recurrences, KV-cache memory balloons linearly, negating the parameter efficiency that made looping attractive. This paper from UC San Diego and Together AI attacks that memory gap head-on. It identifies why naive KV-sharing fails, then introduces a two-stage recipe—stochastic-depth pretraining followed by a targeted post-training correction called BeNice!—that makes KV-sharing reliable across model scales and compatible with other compression tricks.
Why Looped Models Hit a Memory Wall
A middle-looped architecture splits computation into three stages: a prelude block that embeds the input sequence, a recurrent block that iteratively refines a hidden state conditioned on that embedding, and a coda block that produces logits from the final hidden state. During training, the recurrent depth is either sampled stochastically or controlled by an adaptive halting mechanism. At inference, however, a fixed number of recurrences T is chosen, and the model steps through the recurrent block T times.
Each recurrence attends to its own KV-cache. For a context of length L, hidden dimension d, and T recurrences, the KV-cache footprint scales as O(T × L × d). The model weights stay constant because the recurrent block is shared, but the cache grows with every additional pass. Prior work on memory-efficient attention—multi-latent attention (MLA), grouped-query attention (GQA), Kimi delta attention (KDA)—shrinks the per-layer cache size but leaves the linear dependence on depth untouched. Looped models therefore trade parameter efficiency for inference-memory inefficiency.
KV-sharing proposes a simple fix: reuse the final recurrence's KV-cache for all earlier recurrences. If the final cache is close enough to what earlier recurrences would have produced, quality holds while memory drops to that of a single recurrence. Prior work observed this works for some looped models (Huginn, smaller Parcae checkpoints) but fails catastrophically for others (larger Parcae, Ouro) even under identical training setups. The mechanism behind this inconsistency was unexplained.
Fixed-Point Convergence and the Role of Stochastic Depth
The paper's first major contribution is a theoretical and empirical characterization of when KV-shareability emerges. The central insight: KV-shareability is a downstream consequence of fixed-point convergence in the recurrent block. When the recurrent block acts as a contraction mapping on the hidden state, the sequence of hidden states h₁, h₂, …, hₜ converges to a fixed point h*. Because key and value projections are locally Lipschitz in the hidden state, their outputs also converge. The distance between an intermediate cache and the final cache is bounded by the step-to-step residual norm ||hₜ - hₜ₋₁||. As the recurrent block settles, swapping in the final cache perturbs attention outputs by a vanishingly small amount.
Stochastic-depth pretraining induces this contractive behavior. By randomly sampling the recurrence depth per sequence during training—say from a Poisson distribution with mean μ_rec—the model must produce low-loss predictions at many different depths using the same tied weights and a shared readout. This discourages residual drift; successive hidden states cannot diverge because the model might be asked to stop at any depth. Static-depth training lacks this pressure. The paper trains Parcae models from scratch at 140M, 370M, 770M, and 1.3B parameters with both stochastic and static depth. Stochastic-depth models show clean fixed-point convergence in residual norms and attention patterns; static-depth models do not. Crucially, the quality gap between shared and unshared KV evaluation is near zero for stochastic-depth models (ΔCore ≤ 0.25) but large for static-depth models (ΔCore up to 2.00 at 1.3B).
The final recurrence's cache is the optimal broadcast target because it sits closest to the fixed point. Broadcasting an earlier cache would introduce larger substitution error. This explains why KV-sharing works when it works: contractive dynamics align the caches, and the final one is the natural anchor.
When Fixed Points Diverge: Token-Independent Basis Shifts
Fixed-point convergence is necessary but not sufficient. The 770M and 1.3B Parcae models use stochastic depth and converge to fixed points, yet they fail KV-shareability—shared evaluation degrades Core by 3.94 and 7.40 points respectively. The paper diagnoses the failure mode by measuring the mean relative difference between final hidden states under shared versus unshared evaluation. Shareable models (140M, 370M) show differences ≤ 1.69%. Non-shareable models (770M, 1.3B) show differences up to 177.5%.
The failure localizes to the first layer of the first recurrence. Preserving only that layer's KV from the first recurrence (while sharing the rest) restores unshared quality. The error is a token-independent basis shift: the residual perturbation introduced by sharing the final cache into the first recurrence is nearly identical across all tokens. The paper constructs an error subspace from the first-layer output difference and finds average projected pairwise cosine similarity of 0.99—tokens align strongly within this subspace but are orthogonal outside it. This means KV-sharing applies a consistent linear transformation to the residual stream, shifting the fixed point and degrading downstream quality. A naive patching correction—adding the mean projected error back to the residual after the first layer—recovers unshared performance without storing the first recurrence's cache.
BeNice!: A Two-Stage Recipe for Reliable KV-Sharing
The paper converts this analysis into a practical method. Stage 1: pretrain with stochastic depth to induce contractive residual dynamics. Stage 2: post-train with BeNice! to correct the first-recurrence basis shift.
BeNice! defines a hybrid predictor p_θ^(1←T) that runs the first recurrence using the final recurrence's KV-cache. The loss combines standard language modeling on the full unshared T-recurrence pass with an auxiliary term on the hybrid pass:
L_BeNice! = CE(y, p_θ^(T)) + γ · CE(y, p_θ^(1←T))
The cross-entropy on the unshared pass preserves base quality. The auxiliary term trains the first recurrence to produce residuals compatible with the final recurrence's KV-cache. The hyperparameter γ balances the two objectives. During finetuning, the recurrence depth T is sampled from the same distribution used during pretraining. The result: the first recurrence learns to "play nice" with the shared cache, eliminating the basis shift without requiring full KV-sharing rollouts during training (which would be computationally prohibitive).
The method extends naturally to memory-efficient sequence mixers. For multi-latent attention, the final recurrence's latent cache is broadcast. For Kimi delta attention, the final recurrence's state is reused. The same loss formulation applies.
End-to-End Quality: BeNice! Matches Finetuning and Enables Sharing
The paper evaluates on newly trained Parcae models (140M to 1.3B) pretrained on 100B Climbix tokens, then finetuned with BeNice! using three seeds. Table 4 tells the story: without BeNice!, KV-sharing degrades Core by up to 5.58 points (370M) and Core-Extended by up to 3.31 points. With BeNice!, shared Core actually exceeds unshared Core at 140M (+0.14) and 1.3B (+0.04), and the gap shrinks to ≤ 0.39 at 370M and ≤ 0.26 at 770M. Base model quality improves too—BeNice! acts as an effective finetuning recipe regardless of sharing.
Against standard finetuning (same compute, no auxiliary loss), BeNice! matches unshared validation perplexity (Δ ≤ 0.01 PPL) while reducing shared perplexity by up to 9.75 points (45.2% relative). Standard finetuning with KV-sharing suffers 7–33 PPL degradation. BeNice! is the first method to close that gap.
Composition with Memory-Efficient Mixers and Quantization
KV-sharing is orthogonal to other compression axes. The paper tests BeNice! with MLA and KDA sequence mixers. Without BeNice!, both mixers fail KV-shareability (Core drops ~5–10 points). With BeNice!, shared Core degrades by only 0.09 and 0.04 points on average for MLA and KDA respectively. The correction generalizes across attention variants.
For quantization, the paper evaluates fp8, mxfp8, and nvfp4 KV-cache precision. Shared evaluation with quantized caches changes Core by -0.02 to +0.06 points and Core-Extended by +0.02 to +0.08 points on average. ClimbMix perplexity increases just 0.67% on average. The interaction effect between KV-sharing and quantization is ≤ 0.53% in perplexity—essentially additive compression. This means practitioners can stack KV-sharing with low-precision caches for compounding memory savings.
Real-World Inference Gains: Cache Hits, Throughput, TTFT
The paper simulates long-context inference using Mooncake production traces replayed through SGLang. Proxy models are built from GLM 4.7-FP8 by looping the middle transformer blocks. At μ_rec = 8, KV-sharing cuts HBM usage from 63.8 GB to 21.6 GB (66% reduction) and expands token capacity 3.33×. At μ_rec = 16, token capacity grows 6×.
The memory headroom translates to higher cache-hit rates: +8.1 percentage points (31.4% relative) at 22 requests/second. Throughput scales accordingly: +23.6% at 2 req/s. Tail latency drops dramatically—p90 time-to-first-token falls from 139.2 s to 23.2 s (83.6% reduction). Under high concurrency, KV-sharing yields up to 6.7% higher throughput at fixed decode rate, or 34.8% higher median decode rate at fixed throughput. These gains materialize because the fixed KV-cache budget now accommodates more concurrent sequences without eviction and recomputation.
Limitations and Open Questions
The paper acknowledges several constraints. First, BeNice! requires stochastic-depth pretraining; it does not retrofit onto models trained with static depth or ACT-based halting (Ouro). Second, the auxiliary loss introduces a hyperparameter γ that needs tuning per model scale. Third, the analysis assumes middle-looped architectures with implicit halting; full-looped or router-based variants may behave differently. Fourth, while the paper proves KV-cache convergence under contractive dynamics, the proof assumes locally Lipschitz projections—violated by some attention variants with hard non-linearities. Fifth, the inference simulations assume perfect prefix sharing; real workloads with diverse prefixes may see smaller cache-hit improvements.
An intriguing direction the paper opens: can the basis-shift correction be baked into the architecture itself, eliminating the post-training step? The naive patching result suggests a fixed linear correction suffices, hinting at a possible architectural fix (e.g., a learned adapter on the first recurrence's residual). The paper also notes that optimal hyperparameters transfer strongly across recurrence counts, suggesting the contraction strength is predictable.
What This Means for Practitioners
If you are deploying looped models—or considering them for inference-constrained environments—this work changes the calculus. KV-sharing is no longer a brittle heuristic; it is a trainable property. The two-stage recipe is practical: pretrain with stochastic depth (already standard for Parcae and Huginn), then apply BeNice! finetuning (a few thousand steps with an extra loss term). The memory savings are immediate and compose with MLA, KDA, and fp8/mxfn4 quantization. For a 1.3B looped model at 16 recurrences, you can serve 6× more tokens per GPU without quality loss. Tail latency improvements are dramatic enough to shift p90 TTFT from "timeout territory" to "interactive."
The broader lesson: inference-memory gaps in recurrent architectures are tractable when you understand the fixed-point geometry. Contractive dynamics align representations; targeted correction aligns the first step. This pattern may apply beyond KV-caches—any recurrent system where intermediate states are cached and could be shared stands to benefit.