BeNice! 用KV共享为循环架构省下推理内存

循环架构（Looped Architectures）通过将激活反复送入同一组层来提升模型质量，而不增加参数量。但这一设计在推理时有个天然缺陷：每次循环都独立生成KV缓存，内存占用随循环深度线性增长。一篇新论文提出两阶段方案——预训练用随机深度诱导收缩残差，后训练用BeNice!广播KV缓存，让循环模型在保持质量的同时实现KV共享。

BeNice! 论文来源截图

为什么循环架构的KV缓存是个问题

标准Transformer每层只计算一次KV，而循环架构的同一组层会被执行多次。假设循环4次，就等于存了4份KV缓存。上下文越长、循环越深，显存压力越大，直接拖慢长上下文推理的吞吐和首token延迟。

已有研究观察到一种现象：在某些训练好的循环模型里，最后一次循环的KV缓存可以直接给前面几次循环复用，质量几乎不受影响，作者称之为“KV可共享性”（KV-shareability）。问题在于，这个性质并不稳定——即使架构相同、训练方式相同，有些模型具备，有些不具备，背后机制此前并不清楚。

KV缓存（Key-Value Cache）是自回归推理中的一项关键优化：解码每个新 token 时，模型需要attend到此前所有 token 的键（Key）和值（Value）向量，若每次都重新计算，代价会随序列长度平方增长。因此标准做法是把每层已算出的 K、V 张量缓存下来，后续只计算新 token 的部分。这也意味着缓存大小与层数、上下文长度、批大小成正比，是长上下文服务显存占用的主要来源。循环架构把同一组层重复执行多次，若每次循环各自维护一份缓存，等价于把“有效层数”放大了循环倍数，缓存占用随之线性膨胀，这正是本文要解决的瓶颈。

两阶段方法：先让表示收敛，再让KV兼容

论文的核心贡献是弄清了KV可共享性的来源，并给出可复现的实现路径。

第一阶段：随机深度带来收缩残差

预训练阶段引入随机深度（stochastic depth），会诱导出一种收缩性残差动态（contractive residual dynamics）。直观理解是：每一次循环的输出都被拉向一个共同的不动点，不同循环产生的表示越来越接近。表示趋同，KV缓存自然更容易互相替代。这一步在预训练时就把“可共享”的土壤铺好，而不是靠后处理硬掰。

随机深度（Stochastic Depth）最初是作为正则化手段提出的：在训练时以一定概率随机跳过残差分支，使网络在期望意义上等价于大量浅层子网络的集成，从而缓解过拟合、加速训练。其关键结构特征是残差连接——被跳过的层输出直接等于输入。当这种随机跳过与循环结构结合时，残差分支的贡献被反复削弱，多次循环的输出会逐渐向一个共同的不动点收缩，形成所谓收缩性残差动态。直观上，这抑制了循环之间表示的漂移，使不同深度的循环产生高度相似的中间状态，为后续让 KV 缓存跨循环复用提供了表示层面的前提。

第二阶段：BeNice! 广播最终循环的KV

有了第一阶段的基础，作者提出后训练方法BeNice!：把最后一次循环的KV缓存广播到第一次循环，并训练第一次循环产生与之兼容的残差。关键在于，这一过程不牺牲原模型质量——BeNice! 在启用KV共享的同时，验证集困惑度与标准微调（不共享）持平。相比直接拿一个普通训练的循环模型来共享，这个后训练步骤才是稳定触发KV可共享性的开关。

与现有压缩手段互补，而非替代

论文特别验证了KV共享的正交性：它可以和内存高效的序列混合器叠加使用，例如多头潜在注意力（MLA）和Kimi delta attention；也可以配合低精度KV缓存量化，包括fp8、mxfp8、nvfp4。换句话说，BeNice! 不是另一条压缩路线的竞争者，而是能加在已有优化之上的额外一层。对工程团队来说，这意味着它更容易嵌入现有推理栈。

多头潜在注意力（MLA）是 DeepSeek^广告系列模型采用的一种注意力变体：它不直接缓存完整的 K、V 向量，而是先把它们压缩到一个低维潜在空间，推理时再上投影还原，从而大幅削减 KV 缓存体积。Kimi delta attention 则是线性注意力方向上的改进，用增量式的状态更新替代对全序列的显式缓存，把内存复杂度从随长度增长降到接近常数。低精度 KV 量化（fp8、mxfp8、nvfp4）则是从数值表示层面压缩缓存，其中 mxfp8 与 nvfp4 属于带共享缩放因子的微缩放格式，在保持精度的同时进一步降低位宽。本文验证 KV 共享与这些手段正交，意味着它们分别作用于循环维度、序列维度和数值精度维度，可叠加获得复合收益。

模拟长上下文推理的实测收益

在低内存、长上下文的模拟推理场景中，KV共享带来了三项可量化提升：缓存命中率最高提升31.4%，吞吐最高提升23.6%，p90首token时间最多降低83.6%。首token延迟的改善幅度最大，说明KV共享对交互式长上下文应用的体验提升尤其明显——当缓存能跨循环复用，调度器不必反复搬运和重算，排队等待的尾部延迟被显著压缩。

说一下，这些数字来自模拟环境，真实部署中的收益会受具体硬件、批大小和上下文长度分布影响。但论文给出的方向是清晰的：循环架构的内存瓶颈不必靠削减循环次数来换，KV共享提供了一条保持质量的替代路径。

对架构选择的启示

循环架构一直被视为“用小参数换高质量”的性价比方案，推理内存是它落地的主要障碍。BeNice! 的价值在于把这个障碍变成可训练、可复现的工程问题：预训练阶段用随机深度埋下收敛性，后训练阶段用广播加兼容训练完成收口。对正在评估循环架构的团队，这套方法值得纳入推理优化清单，尤其是那些已经用上MLA或低精度KV量化的长上下文服务。