BeNice! Uses KV Sharing to Save Inference Memory for Looped Architectures

Looped architectures improve model quality by feeding activations repeatedly through the same set of layers without increasing parameter count. But this design has a natural drawback at inference time: each loop generates its own KV cache independently, so memory usage grows linearly with loop depth. A new paper proposes a two-stage approach—using stochastic depth during pretraining to induce contractive residuals, then broadcasting the KV cache post-training with BeNice!—enabling looped models to share KV caches while maintaining quality.

Why KV Caches Are a Problem for Looped Architectures

A standard Transformer computes KV once per layer, whereas a looped architecture executes the same set of layers multiple times. With 4 loops, that means storing 4 copies of the KV cache. The longer the context and the deeper the loop, the greater the memory pressure—directly slowing down long-context inference throughput and time-to-first-token.

Prior research has observed a phenomenon: in some trained looped models, the KV cache from the final loop can be directly reused by earlier loops with almost no quality degradation, which the authors call "KV-shareability." The problem is that this property is unstable—even with identical architectures and training methods, some models have it and some don't, and the underlying mechanism was previously unclear.

The KV cache (Key-Value Cache) is a key optimization in autoregressive inference: when decoding each new token, the model needs to attend to the key (Key) and value (Value) vectors of all previous tokens. If recomputed each time, the cost grows quadratically with sequence length. The standard approach is therefore to cache the computed K and V tensors for each layer, computing only the new token's portion subsequently. This also means cache size is proportional to the number of layers, context length, and batch size, making it the primary source of memory usage in long-context serving. Looped architectures execute the same set of layers multiple times—if each loop maintains its own cache, this is equivalent to multiplying the "effective layer count" by the loop factor, causing cache usage to inflate linearly. This is precisely the bottleneck this paper addresses.

Two-Stage Method: First Converge Representations, Then Make KV Compatible

The paper's core contribution is clarifying the source of KV-shareability and providing a reproducible implementation path.

Stage One: Stochastic Depth Induces Contractive Residuals

Introducing stochastic depth during pretraining induces contractive residual dynamics. Intuitively: each loop's output is pulled toward a common fixed point, making representations from different loops increasingly similar. When representations converge, KV caches naturally become more interchangeable. This step lays the groundwork for "shareability" during pretraining, rather than forcing it through post-processing.

Stochastic depth was originally proposed as a regularization technique: during training, residual branches are randomly skipped with a certain probability, making the network equivalent in expectation to an ensemble of many shallower subnetworks, thereby mitigating overfitting and accelerating training. Its key structural feature is the residual connection—the output of a skipped layer equals its input directly. When this random skipping combines with looped structures, the contribution of residual branches is repeatedly weakened, and outputs from multiple loops gradually contract toward a common fixed point, forming what is called contractive residual dynamics. Intuitively, this suppresses representation drift between loops, causing loops at different depths to produce highly similar intermediate states, providing the representational prerequisite for subsequently reusing KV caches across loops.

Stage Two: BeNice! Broadcasts the Final Loop's KV

With the foundation from stage one, the authors propose the post-training method BeNice!: broadcasting the final loop's KV cache to the first loop and training the first loop to produce compatible residuals. Crucially, this process does not sacrifice the original model's quality—BeNice! achieves validation perplexity on par with standard fine-tuning (without sharing) while enabling KV sharing. Compared to directly sharing with a conventionally trained looped model, this post-training step is the switch that reliably triggers KV-shareability.

Complementary to Existing Compression Methods, Not a Replacement

The paper specifically verifies the orthogonality of KV sharing: it can be stacked with memory-efficient sequence mixers such as multi-head latent attention (MLA) and Kimi delta attention; it can also be combined with low-precision KV cache quantization, including fp8, mxfp8, and nvfp4. In other words, BeNice! is not a competitor to other compression approaches but an additional layer that can be added on top of existing optimizations. For engineering teams, this means it's easier to integrate into existing inference stacks.

Multi-head latent attention (MLA) is an attention variant adopted by the DeepSeek^Ad series of models: instead of directly caching full K, V vectors, it first compresses them into a low-dimensional latent space and projects them back during inference, substantially reducing KV cache size. Kimi delta attention is an improvement in the linear attention direction, replacing explicit caching of the full sequence with incremental state updates, reducing memory complexity from growing with length to near-constant. Low-precision KV quantization (fp8, mxfp8, nvfp4) compresses the cache at the numerical representation level, where mxfp8 and nvfp4 are micro-scaling formats with shared scaling factors that further reduce bit width while maintaining precision. This paper verifies that KV sharing is orthogonal to these methods, meaning they act on the loop dimension, sequence dimension, and numerical precision dimension respectively, and can be combined for compound benefits.

Measured Gains in Simulated Long-Context Inference

In simulated inference scenarios with low memory and long contexts, KV sharing delivers three quantifiable improvements: cache hit rate up to 31.4% higher, throughput up to 23.6% higher, and p90 time-to-first-token reduced by up to 83.6%. The improvement in first-token latency is the largest, indicating that KV sharing particularly enhances the experience of interactive long-context applications—when caches can be reused across loops, the scheduler doesn't need to repeatedly move and recompute data, significantly compressing tail latency from queued requests.

Note that these numbers come from simulated environments; real-world deployment gains will be affected by specific hardware, batch size, and context length distribution. But the paper's direction is clear: the memory bottleneck of looped architectures doesn't have to be traded off by reducing loop count—KV sharing offers an alternative path that maintains quality.

Implications for Architecture Selection

Looped architectures have long been viewed as a cost-effective approach to "trading small parameters for high quality," with inference memory being the main obstacle to deployment. BeNice!'s value lies in turning this obstacle into a trainable, reproducible engineering problem: using stochastic depth during pretraining to establish convergence, and broadcasting plus compatible training during post-training to close the loop. For teams evaluating looped architectures, this method is worth adding to the inference optimization checklist—especially for long-context services that already use MLA or low-precision KV quantization.

One-line Summary

A two-stage method enables looped architectures to share KV caches, significantly improving long-context inference throughput and time-to-first-token.

Paragraph Summary

Looped architectures improve quality by repeatedly executing the same set of layers without increasing parameters, but each loop independently generates a KV cache, causing memory to grow linearly with loop depth and becoming the primary bottleneck for long-context inference. The paper reveals the mechanism behind why some looped models exhibit "KV-shareability" while others don't, and proposes a two-stage approach: during pretraining, stochastic depth is introduced to induce contractive residual dynamics that converge representations across loops; during post-training, BeNice! broadcasts the final loop's KV cache to the first loop and trains it to produce compatible residuals, matching standard fine-tuning quality. The method is orthogonal to and stackable with MLA, Kimi delta attention, and fp8/mxfp8/nvfp4 quantization. In simulated long-context scenarios, cache hit rate improves by up to 31.4%, throughput by 23.6%, and p90 time-to-first-token drops by 83.6%.