一种自剪枝 Transformer：借助 Universal Attention 实现极致 KV-cache 压缩

arXiv:2610.09051v1 公告类型：新增 摘要：现代 LLM 的庞大 KV-cache 规模成为高效部署的障碍。近期研究探索以基于衰减的替代机制取代注意力层中的 RoPE 位置嵌入，从而在推理期间对 KV-cache 进行剪枝。然而，这些衰减函数的表达能力有限，在实践中会退化为类似滑动窗口的驱逐模式。在本工作中，我们提出了一个统一框架，用于实现互补且新颖的衰减机制，在保留具有表达力的 RoPE 嵌入和 Softmax 注意力的同时，捕捉复杂的键统计特征及其交互作用。由此得到的 Universal Attention 是一种具有高度表达力且可端到端训练的架构，其复合衰减机制可作为一种自然的、自适应的剪枝标准，移除对注意力计算贡献最小的 token。实验表明，在自然语言和合成任务数据上，Universal Attention 实现了最先进的 10 倍压缩效果，同时相较于最先进的基线方法以及未剪枝的 oracle，进一步提升了下游性能。此外，在长度为 16k 的场景下，它以前所未有的 25 倍压缩展现出更优越的长上下文泛化能力。