Hayden Prairie*, Jacob Ryan*, Taylor Berg-Kirkpatrick, and Dan Fu

Paper, GitHub, HuggingFace

* Joint Authorship

TL;DR: We present BeNice!, a post-training method that enables KV-shareability in looped architectures, allowing a single KV-cache to be used across all recurrences without quality degradation. KV-shareability is complementary to other KV-cache compression methods, composing with efficient sequence-mixers and KV-quantization, and enables memory-efficient, higher-throughput inference for looped models.

A KV-cache for every recurrence sounds expensive. Do we need them all?

Looped architectures have entered the spotlight for language modeling, with several rumors that GPT-Astra and Gemini are looped models. In our last work Parcae, we showed that looping offers a new, efficient avenue to scale quality while keeping parameter count fixed. However, as we scale recurrence and depth, KV-cache size still grows linearly. We wonder: Can we scale recurrences while keeping cache size constant?

Exploring this question on top of our prior work Parcae, we have found that we can directly share the final-recurrence KV-cache across all prior recurrences with near-lossless quality degradation, a phenomenon we term KV-shareability. Unfortunately, the mechanisms underlying KV-shareability remain poorly understood, with models trained under identical setups differing in whether they exhibit KV-shareability. We study the emergence of KV-shareability, and propose a two-stage pipeline to enable it: stochastic depth pretraining and BeNice! post training. Together, these stages:

• Achieve near-lossless KV-sharing while maintaining base quality. We find that BeNice! maintains base model quality without sharing, while maintaining Core-Ext within 0.06 points and validation perplexity within 0.01 when using KV-sharing.
• Enables KV-shareability as a complementary axis of KV compression. We find that BeNice! is complementary to other KV compression techniques, composing with MLA and KDA sequence mixers and up to 4-bit KV-quantization.
• Accelerates throughput and reduces time-to-first-token (TTFT) in memory-constrained settings. The memory savings of KV-shareability directly translate to faster inference, with up to 23.6% more throughput and an 83.6% reduction in p90 TTFT for memory-constrained settings.

Naively sharing across recurrences works … sometimes

As state-of-the-art models max out the high-bandwidth (HBM) memory of frontier chips and more deployments moving to local machines, the need to scale quality without scaling memory footprint has increased. While looped models keep parameter overhead constant through iteratively looping activations over a block of weights, they linearly increase the KV-cache footprint with each extended recurrence.

Interestingly, some looped architectures have a unique property to directly share the final recurrence’s KV-cache across all prior recurrences, without degrading quality, which we show in the table below. We use the term KV-sharing to indicate the process of reusing a cache across multiple recurrences, while reserving the term KV-shareability for models that do so without quality degradation.

So what enables KV-shareability in the first place? Our key observation is that contractive residual dynamics (i.e., convergence to a fixed-point) drives residual states towards fixed points, whose agreement under shared and unshared KV-caches enables KV-shareability. Put simply, the internal trajectory taken with and without KV-sharing ends up at the same place even if they take different paths.

Enabling a contractive recurrent block

To enforce contractive residual dynamics, we train with stochastic depth, i.e., randomly sampling recurrence depth throughout training. Intuitively, random sampling depth means each recurrence must produce a low-loss prediction regardless of depth, discouraging drift and encouraging a contractive residual.

Now, why are contractive residual dynamics so important? Consider a looped architecture in which the recurrent block $\mathcal{R}$ learns a contraction upon the residual hidden state $h_t$, such that $h_t \to \hat h$ where $\hat h$ is a fixed point. We see this phenomenon appearing in all of our prior Parcae architectures, as seen below.

By swapping in a converged KV-cache $(\hat k, \hat v)$ to a prior recurrence $(k_t,v_t)$, it perturbs the output by a term bounded by the residual difference $\lVert h_t - \hat h \rVert$ (see paper for additional details). This is necessary for enabling KV-shareability. To test this empirically, we train Parcae models with stochastic depth (i.e., randomly sampling depth during training) and static depth (i.e., holding depth constant during training). As shown in the table below, models trained with stochastic depth have KV-shareability while learning a contraction on the residual (see the figure above). Alternatively, models trained with static depth have no incentive to keep their hidden state near a common limit (see the figure above), which prevents them from learning a shareable representation across recurrences (see the table below).

KV-shareability suggests that looped architectures should favor stochastic depth training over methods such as ACT (e.g., Ouro-based architectures). By inducing a contraction on the residual, stochastic depth training enables KV-shareability without explicitly optimizing for it.

But why does KV-shareability sometimes fail with stochastic depth?

Okay, but stochastic depth is clearly not enough. We can see that KV-shareability does not emerge in all models trained with stochastic depth (e.g., Parcae 770M and 1.3B). While stochastic depth learns a fixed point across a model’s residual, we find that it sometimes is not fully sufficient to converge to the same fixed point with and without KV-sharing. So what gives?

Digging into the failure point of KV-sharing, we find that models lacking KV-shareability introduce a token-independent basis shift in the first recurrence, causing the model to converge to a different fixed point. The fixed point failure can be visualized in the figure below: models with KV-shareability (e.g., Parcae 140M and 370M) converge to the same fixed points, while models without KV-shareability (e.g., Parcae 770M and 1.3B) converge to different fixed points. We confirm this numerically, where we find that the average relative tokenwise difference with and without KV-sharing is 1.69% for Parcae 370M and 177.5% for Parcae 1.3B.

Now picking through the internals of the model, we find that we can regain quality by allocating a separate KV-cache for the first layer during the first recurrence, and sharing all future KV-caches (see the table below). Thus first layer of the first recurrence introduces failures under KV-sharing, which we find to be a token-independent basis shift: by constructing a subspace $B$ from the error delta between reusing the final KV-cache and the recurrence-correct KV-cache and projecting each token into $B$, we observe an average projected pairwise cosine similarity of 0.99 (see figure below).

Thus, the basis shift introduced by KV-cache reuse in the first layer of the first recurrence is exactly what alters the converged representation at later recurrences, degrading quality.

BeNice! Sharing is Caring in Looped Models

Okay, okay, so now that we know what’s happening, how can we enable KV-shareability? We propose a two-stage pipeline: first, we pretrain with stochastic depth to enable contractive recurrent dynamics in the residual; second, we finetune with BeNice!, a post-training algorithm that directly trains the first recurrence pass to be compatible with future KV-caches, correcting KV-sharing failures. Specifically, consider the ordinary $T$-recurrence predictor $p_\theta^{(T)}$ and a hybrid predictor $p_\theta^{(1\leftarrow T)}$ that uses first-recurrence queries and residuals with final-recurrence keys and values.

During finetuning, we sample $T$ from the same recurrence-depth distribution $\Lambda$ used during pretraining. Intuitively, the auxiliary term trains the first recurrence to be compatible with future recurrences’ KV-caches, while the standard language-modeling term maintains base quality at later depths. For more details on our two-stage approach and for generalization of BeNice! to memory-efficient sequence mixers, please check out the paper.

Back to Language Modeling

Now the results. Not only does our pipeline enable KV-shareability, it also achieves comparable quality to standard finetuning during unshared decoding.

To validate BeNice!, we pretrain new Parcae models from scratch with stochastic depth. We find that all models suffer from the same KV-cache reuse failure points as above, and that BeNice! corrects these failures in all settings, enabling KV-shareability.

Furthermore, using prior Parcae checkpoints, we perform a small amount of post-training using both BeNice! and standard finetuning. We find that BeNice! directly matches the quality of standard fine-tuning without KV-cache reuse, while also achieving KV-shareability.

KV-shareability Complements Other KV Compression Techniques

Now that we can reliably enable KV-shareability, can it be composed with other KV compression techniques? KV cache overhead can generally be expressed by the product of its numerical precision, the per-layer overhead, and the depth of the model.

To test per-layer overhead, we train models from scratch with memory-efficient sequence mixers such as multi-latent and kimi-delta attention. As seen in the figure above, we find that BeNice! is able to effectively enable KV-shareability in all models.

To test numerical precision, we downsample the numerical precision of KV-cache entries with and without KV-cache reuse. Again, we observe comparable results using lower precision with and without KV-cache reuse, as seen in the figure above. Thus, KV-shareability persists with KV-quantization. Overall, KV-shareability is compatible with KV-quantization and memory-efficient sequence-mixers.

Let’s use the saved space to accelerate inference!

Finally, let’s turn that saved space directly into real inference savings. We build a synthetic environment to explore how KV-shareability impacts real inference settings (e.g., with real request traces, synthetic model scales, and optimized inference serving platforms). A more detailed look at how we built this can be found in the paper.

Not only does KV-shareability save memory, but those savings directly increase the total tokens resident on-device, which can improve cache-hit rates in memory-constrained environments. These savings can be seen in the figure above.

Higher cache-hit rates directly translate into better overall throughput and lower TTFT. At low loads (e.g., 0.5 req/s), KV-shareability matches the unshared baseline in TTFT and throughput; however, at high loads (e.g., 2 req/s), KV-shareability improves throughput by 23.6% and reduces p90 TTFT from 139.2s to 23.2s. When simulating the concurrency frontier (e.g., with increasing number of concurrent users), we see that KV-shareability provides up to 6.7% higher throughput at a fixed median decode rate and 34.8% higher median decode rate at a fixed throughput.

What’s next & trying it out yourself

We are really excited to take looped models to the next level, pushing parameter efficiency to its limits and developing new inference algorithms to exploit that efficiency. We are releasing all of our code and models to further help the process. We will be releasing more work soon, further building off these ideas, so keep an eye out for what comes next!

If you have any questions or want to help push forward the frontier of looped architectures, please feel free to reach out to Hayden Prairie (hprairie@ucsd.edu) or Jacob Ryan (jmryan@ucsd.edu).

Acknowledgements

We would like to thank Together AI for collaborating with us and providing compute for these experiments. We graciously acknowledge the support of AMD through the FAR research grant. This research was supported in part by the NSF under grants 2146151 and 2422214.

We would also like to thank Hansen Lillemark for his helpful feedback on this blog post.