Paper Souce: Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (arxiv.org)(ICLR’2024)

Observations

Observation 1-1

Observation 1-2

对于一个大模型，同一层中的不同head，不同层的不同head之间，它们学到的注意力机制是不一样的。有的heads在计算注意力时，更加关注special tokens(<s>)；有的更加关注标点（punctuations），有的关注局部性（locality，recent tokens）；有的关注高频 tokens；有的全局都关注（all tokens）。

所以，本文提出针对这5类heads实行不同的 KV Cache 压缩策略

Observation 2

那么在生成过程中，各个head的注意力关注点是否会发生转变呢？

上图表明，用户提供的prompt在生成过程中各个head的注意力结构具有一致性（consistent），基于此，假定已确定的注意力模式在未来的生成步骤中将保持不变，因此在只需要根据prompt进行一次模型剖析（model profiling，即分析head的注意力结构并确定压缩策略）。

如，Layer 33 Head 0 与 Layer 23 Head 2 几乎只关注 special tokens；
Layer 23 Head 0中，标点和局部tokens也有较大比重；
Layer 23 Head 3则有超过10%的注意力分数给到Others部分，意味着它不适合压缩KV Cache（它关注全局tokens）

此外，可以看到大部分注意力都会给到 special tokens，这证明后文提到的混合策略中的贪心方法的合理性。

Method： Adaptive KV Cache Compression

如前文观察到的，大模型中的注意力heads通常所学习到的注意力机制有所差异，这表明需要根据不同的head去“定制”各自的压缩策略。

  1. 在预填充阶段，为了给各个head选择最合适的压缩策略，需要进行模型剖析（model profiling）来确定各个注意力head的注意力模式。

    • Model Profiling：

      Equation 1

      • Compression Policies：

        • Special Tokens: 只保留起始符<s>，指令符[INST]等special tokens的KV Cache。CspecialC_{special}
        • Punctuation: 只保留 .,,,:,?等标点符号的KV Cache。CpunctC_{punct}
        • Locality:若某个上下文token与当前token的相对距离超过一定阈值时则去除该token的KV Cache。ClocalC_{local}
        • Frequency(Heavy Hitter): 记录每个 token 的累计注意力分数，将这些分数称为 toekn frequency, 只保留 分数最高的几个token（也就是“最重要的几个token”）。CfrequentC_{frequent}

      • Hybrid Policies:

        混合策略中采用贪心方法，得到如下策略集：

        Hybrid Policies

        这里需要解释为什么以CspecialC_{special}作为各个策略的一部分：

        • 如前所述，大部分较高的注意力分数都集中在 special tokens，因此CspecialC_{special}对于保留模型精度至关重要
        • 一句话中的 special tokens 通常不超过5个，因此对于cache大小影响较小（这点同样适用于CpunctC_{punct}）

        Algorithm 1

  2. 在生成阶段，基于所选的压缩策略来管理各个 token 的 KV Cache

    Algorithm 2