BeNice!：通过 KV 共享弥合循环架构的推理内存差距

arXiv:2610.31633v1 公告类型：新论文 摘要：循环架构通过让激活值反复经过同一组网络层，在保持参数量不变的同时提升模型质量。在推理过程中，每次循环都会生成独立的 KV 缓存，导致内存占用线性增长。为应对这一挑战，已有研究提出了 KV 共享（即跨深度共享 KV 缓存），并发现，在某些情况下，可以将最后一次循环的 KV 缓存与之前的循环共享，而不损失模型质量。我们将这一性质称为 KV 可共享性。然而，KV 可共享性的内在机制仍未得到充分理解，即便采用相同方式训练的循环架构，也并不总能表现出这一性质。本文研究如何使循环模型具备 KV 可共享性。我们提出了一种两阶段方法：第一阶段，在预训练期间采用随机深度，从而诱导出收缩性残差动力学，推动不同循环中的表征趋向共同极限。第二阶段，我们提出后训练方法 BeNice!，将最后一次循环的 KV 缓存广播至第一次循环，训练其产生兼容的残差，同时保留原始模型的质量。因此，BeNice! 在实现 KV 可共享性的同时，能够达到标准微调在不共享 KV 缓存时的验证集困惑度水平。利用 BeNice!，我们证明，KV 共享与其他 KV 压缩方法具有互补性，可与内存高效的序列混合器（例如多潜变量注意力、Kimi delta attention）以及低精度 KV 缓存量化（例如 fp8、mxfp8、nvfp4）结合使用。最后，在模拟的低内存、长上下文推理场景中，KV 共享将缓存命中率最多提升 31.4%，将吞吐量最多提升 23.6%，并将首个 token 生成时间的第 90 百分位数（p90）最多降低 83.6%。