YOCO++：通过KV残差连接提升大模型推理效率

本文提出了YOCO++，通过在跨层KV压缩方法中引入加权残差连接，在不牺牲训练和推理效率的前提下，显著提升了模型性能，实现了50% KV缓存压缩率下的最先进表现。

大语言模型的推理效率一直是制约其大规模部署的关键瓶颈。随着模型规模的增长，注意力机制中的键值（KV）缓存占用的内存呈线性增长，这不仅限制了可处理的序列长度，也增加了推理成本。

跨层KV压缩技术应运而生，其核心思想是在不同层之间共享KV缓存，从而大幅减少内存占用。然而，这种压缩通常伴随着性能下降——模型容量受限，表达能力减弱。如何在压缩率和模型性能之间取得平衡，成为这一领域的核心挑战。

YOCO++的提出，正是为了解决这一难题。通过在YOCO框架基础上引入KV残差连接，新方法在不增加计算开销的情况下，显著提升了压缩后模型的性能。

YOCO（You Only Cache Once）是一种创新的跨层KV压缩方法。其核心设计是将Transformer层分为两半：底层（bottom-half）和顶层（top-half）。底层每层独立计算自己的KV，而顶层则共享中间层的KV缓存。

这种设计基于一个关键观察：在Transformer中，不同层对KV信息的需求存在差异。底层需要精细的局部特征提取，而顶层更多地关注全局语义整合。因此，让顶层共享KV是一种合理的权衡。

然而，YOCO的简单共享策略也存在局限。顶层完全依赖中间层的KV，失去了直接接触底层信息的机会，这可能导致信息瓶颈和表达能力下降。

YOCO++的核心改进是在每层底层和底层之间引入加权的KV残差连接。具体来说：

对于每个底层，我们不仅使用它自己计算的KV，还通过残差连接引入底层（最底层）的KV信息。这种设计允许信息从底层直接流向顶层，缓解了YOCO中的信息瓶颈问题。

残差连接的权重是可学习的参数，模型可以根据任务需求自动调整不同层KV信息的混合比例。这种灵活性使得YOCO++能够在保持压缩效率的同时，恢复部分因压缩而损失的表达能力。

YOCO++的残差连接实现简洁而高效。对于第i层底层（i从1到N/2，N为总层数），其KV计算如下：

KV_i = KV_i_base + α_i * KV_bottom

其中，KV_i_base是该层原本计算的KV，KV_bottom是最底层的KV，α_i是可学习的权重系数。

这种设计的优势在于：

计算开销极小：残差连接只涉及简单的向量加法，计算成本可以忽略不计。

参数增量有限：每层只需学习一个标量权重，对模型总参数量影响微乎其微。

训练稳定性：残差连接有助于梯度流动，可能改善深层网络的训练稳定性。

实验结果表明，YOCO++在50% KV缓存压缩率下达到了跨层KV压缩方法的最先进性能。具体而言：

相比原始YOCO，YOCO++在多个基准测试上都取得了显著提升。这种提升在各种任务类型上都有体现，包括语言建模、阅读理解、常识推理等。

更重要的是，YOCO++有时甚至超越了未压缩的标准Transformer。这表明，通过精心设计的压缩策略，我们不仅可以节省内存，还可能获得更好的泛化性能——可能是残差连接起到了某种正则化作用。

与其他的KV缓存压缩方法相比，YOCO++在相同压缩率下 consistently 取得更好的结果。这包括基于量化的方法、基于稀疏化的方法，以及其他的跨层共享方法。