BeNice! Closing Looping's Inference-Memory Gap via KV-Sharing
Looped language models reuse the same parameter block many times to improve quality without inflating weight memory, but each recurrence still stores its own KV-cache, so inference memory grows linearly with depth. This paper studies when and why the final recurrence's KV-cache can be shared across earlier recurrences—a property the authors call KV-shareability—finding that contractive residual dynamics are necessary but not sufficient. They introduce a two-stage recipe (stochastic-depth pretraining plus a lightweight post-training corrector called BeNice!) that reliably enables near-lossless KV-sharing and composes with existing cache-compression techniques.
The paper offers a compelling mechanistic story and a practical training recipe for a real inference-memory bottleneck in looped models. The analysis linking fixed-point convergence to KV-shareability is intuitively appealing, and the BeNice! auxiliary-loss fix is simple yet effective in the reported regime. That said, the work is evaluated only on small-to-medium Parcae-style models (up to 1.3B), and the two-stage requirement limits its applicability to already-trained checkpoints that were not pretrained with stochastic depth.
The strongest element is the careful empirical diagnosis of why KV-sharing sometimes fails. The authors do not simply propose a loss and show metrics; they first establish that stochastic depth induces contractive dynamics, identify that convergence to a fixed point is necessary, and then isolate a token-independent basis shift in the first recurrence as the specific failure mode. The proposed BeNice! loss is elegant, jointly optimizing standard depth and a hybrid first-recurrence/shared-KV path with a single hyperparameter γ. The inference simulations in Section 5.3 go beyond standard perplexity benchmarks, showing that KV-shareability improves cache hit rates by up to 31.4%, throughput by up to 23.6%, and p90 TTFT by up to 83.6% under realistic memory pressure.
The evaluation is limited to one architectural family (Parcae-style middle-looped models) and never exceeds 1.3B parameters, leaving open whether the contractive-dynamics story and BeNice! correction scale to larger models or to fully-looped models like Ouro. The theoretical analysis assumes locally-Lipschitz projections and contraction, but only verifies these empirically via relative residual change; no spectral or Lipschitz estimates are reported for the trained models. In addition, even with stochastic depth, KV-shareability is not guaranteed: the paper notes that models such as Parcae 770M and 1.3B converge to fixed points yet still fail to share caches, implying that the mechanism is not fully captured by contraction alone. Finally, BeNice! is a post-training correction that still benefits from a specific pretraining recipe; it cannot be directly applied to arbitrary looped checkpoints.
The evidence supports the core claims within the tested regime. Table 2 cleanly isolates stochastic depth as the key pretraining factor enabling shareability, and Table 5 shows that BeNice! matches standard finetuning in the unshared regime while closing the shared-cache quality gap (validation perplexity improves by 9.75 points over naively shared baselines). The composition experiments with MLA, KDA, and low-precision quantization are valuable and demonstrate orthogonality. Comparisons to prior work are fair: the authors acknowledge that Huginn exhibited KV-shareability out-of-the-box and that Ouro degrades severely when sharing. However, the simulated inference results rely on proxy models constructed by looping middle blocks of open-weights checkpoints, which adds a layer of indirection between the quality benchmarks and the latency claims.
The paper provides extensive architectural and hyperparameter details in Appendix N and Appendix O, including model configs, optimizer settings, and stochastic-depth distributions. However, no public release of code, checkpoints, or training data splits is mentioned in the main text, which would significantly block independent reproduction. Reproducing the KV-shareability perplexity evaluations is particularly costly because shared-cache inference requires token-by-token rollout at O(T|s|^2) rather than a single parallel forward pass. The inference workload simulations further depend on Mooncake traces, the sglang framework, and specific proxy-model construction; without releasing these artifacts, the exact throughput and TTFT numbers would be difficult to replicate.
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%.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.