A2SF: Accumulative Attention Scoring with Forgetting Factor for Token Pruning in Transformer Decoder
作者: Hyun-rae Jo, Dongkun Shin
分类: cs.CL, cs.LG
发布日期: 2024-07-30 (更新: 2024-07-31)
备注: 11 pages(9 pages + reference 2 pages), 6 figures
💡 一句话要点
提出A2SF，通过引入遗忘因子解决Transformer解码器中token剪枝的累积注意力评分偏差问题。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Transformer解码器 KV缓存压缩 Token剪枝 累积注意力评分 遗忘因子 长序列处理 大型语言模型
📋 核心要点
- 现有基于累积注意力评分的token剪枝方法在Transformer解码器中存在偏差，导致token重要性评估不准确。
- A2SF通过引入遗忘因子，对旧token的注意力评分进行惩罚，从而实现token间更公平的重要性比较。
- 实验表明，A2SF在OPT和LLaMA模型上提高了准确性，尤其在LLaMA 2上，1-shot和0-shot分别提升高达7.8%和5.1%。
📝 摘要（中文）
Transformer架构的大型语言模型（LLM）由于KV缓存而面临内存瓶颈问题，尤其是在处理长序列时。先前的研究提出了KV缓存压缩技术，该技术基于累积注意力评分识别不重要的token，并从KV缓存中移除它们，因为只有少数token在注意力操作中起重要作用。然而，我们观察到现有的累积注意力评分不适用于Transformer解码器结构。在解码器模型中，由于掩码的影响，注意力评分的累积次数因token出现的顺序而异，导致token之间的比较不公平。为了解决这个问题，我们提出了带有遗忘因子的累积注意力评分（A2SF）技术，该技术在注意力评分累积过程中引入了遗忘因子。A2SF通过随时间重复将遗忘因子乘以注意力评分来惩罚来自旧token的过去注意力评分。因此，较旧的token会受到更大的惩罚，从而在不同年龄的token之间提供公平性。通过token之间的公平比较，我们可以更有效地选择重要的token。我们已通过OPT和LLaMA模型中的A2SF验证了准确性的提高，并且A2SF在1-shot和0-shot设置下将LLaMA 2的准确性分别提高了高达7.8％和5.1％。
🔬 方法详解
问题定义：论文旨在解决Transformer解码器中，由于mask机制导致的累积注意力评分（Accumulative Attention Score）在不同位置token间的不公平比较问题。现有方法直接累加注意力权重，忽略了token出现顺序的影响，使得靠前的token更容易获得更高的累积评分，从而影响了token重要性的准确评估，最终导致剪枝效果不佳。
核心思路：核心思路是在累积注意力评分的过程中引入一个“遗忘因子”（Forgetting Factor）。该因子会随着token出现的时间推移，不断衰减其历史注意力评分的贡献。这样，即使是较早出现的token，其累积评分也会因为遗忘因子的作用而降低，从而与其他位置的token进行更公平的比较。
技术框架：A2SF方法主要是在Transformer解码器的注意力计算过程中进行改进。具体来说，在计算每个token的累积注意力评分时，不是简单地将所有历史注意力权重相加，而是将之前的注意力权重乘以一个小于1的遗忘因子。这个过程会重复进行，使得越早出现的token的注意力权重衰减得越多。最终的累积注意力评分将用于token重要性的排序和剪枝。
关键创新：关键创新在于引入了遗忘因子，解决了Transformer解码器中累积注意力评分的偏差问题。与现有方法相比，A2SF能够更公平地评估不同位置token的重要性，从而实现更有效的token剪枝。这种方法简单有效，易于集成到现有的Transformer模型中。
关键设计：遗忘因子是一个小于1的超参数，控制着历史注意力权重衰减的速度。论文中可能需要实验确定合适的遗忘因子取值。此外，A2SF可以与不同的token剪枝策略结合使用，例如，可以根据A2SF计算出的评分，选择保留评分最高的top-k个token。
🖼️ 关键图片
📊 实验亮点
实验结果表明，A2SF在OPT和LLaMA模型上均取得了显著的准确性提升。特别是在LLaMA 2模型上，A2SF在1-shot和0-shot设置下分别实现了高达7.8%和5.1%的准确率提升。这些结果验证了A2SF在Transformer解码器中进行token剪枝的有效性。
🎯 应用场景
A2SF技术可应用于各种基于Transformer解码器的语言模型，尤其是在资源受限的场景下，例如移动设备或边缘计算。通过减少KV缓存的大小，A2SF可以降低内存占用，提高推理速度，从而使得大型语言模型能够在这些平台上高效运行。此外，该技术也有助于提升长文本处理能力，例如机器翻译、文本摘要等。
📄 摘要（原文）
Recently, large language models (LLM) based on transformers are facing memory bottleneck issues due to KV cache, especially in long sequence handling. Previous researches proposed KV cache compression techniques that identify insignificant tokens based on Accumulative Attention Scores and removes their items from KV cache, noting that only few tokens play an important role in attention operations. However, we have observed that the existing Accumulative Attention Score is not suitable for the transformer decoder structure. In the decoder model, the number of times the Attention Score accumulates varies depending on the order of token appearance due to the effect of masking, causing an uneven comparison between tokens. To solve this, we propose Accumulative Attention Score with Forgetting Factor (A2SF) technique, which introduces a Forgetting Factor in the Attention Score accumulation process. A2SF applies a penalty to the past Attention Score generated from old tokens by repeatedly multiplying the Forgetting Factor to the Attention Score over time. Therefore, older tokens receive a larger penalty, providing fairness among different ages of tokens. Through the fair comparison among tokens, we can more effectively select important tokens. We have verified the accuracy improvement through A2SF in the OPT and LLaMA models and A2SF improves the accuracy of LLaMA 2 by up to 7.8% and 5.1% on 1-shot and 0-shot.