[论文] BeNice! Closing Looping's Inference-Memory Gap via KV-Sharing

论文概要

研究领域: ML 作者: Hayden Prairie, Jacob M. Ryan, Taylor Berg-Kirkpatrick, Daniel Y. Fu 发布时间: 2026-10-22 arXiv: 2610.31633

中文摘要

循环架构通过反复传递激活通过同一块层来在固定参数量下扩展质量，但每次递归产生独立KV缓存导致内存线性增长。本文研究如何实现KV可共享性（最终递归的KV缓存可与早期递归共享而不损失质量）。提出两阶段方法：预训练中使用随机深度诱导收缩残差动力学，使表示跨递归趋向共同极限；后训练阶段提出BeNice!，将最终递归的KV缓存广播到首次递归，训练其产生兼容残差。BeNice!使KV共享与多种KV压缩方法互补组合，在模拟低内存长上下文推理中，缓存命中率提升31.4%，吞吐量提升23.6%，p90首token时间降低83.6%。

原文摘要

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, 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. 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 t...

#论文 #arXiv #ML #小凯