Palu: Compressing KV-Cache with Low-Rank Projection

202407，第一作者来自阳明交通大学

问题引入：

• MQA/GQA/MLA这些压缩方式，要么折损性能，要么必须在预训练阶段就做；相比之下，训练后 KV-Cache 压缩技术，比如量化和token驱逐。

• 量化方法旨在减少表示每份数据所需的位宽，而令牌淘汰技术则专注于保留部分 KV-Cache。然而，这两种方法都忽略了 KV-Cache 中隐藏的维度，其中往往存在大量的冗余。

• 引入了 Palu，这是一个基于低秩投影的 KV-Cache 压缩框架，它通过降低 KV 张量的隐藏维度，为现有的量化方法和令牌淘汰方法提供了一个额外的正交压缩维度。

• 如果直接将缓存矩阵映射到低秩空间，在运行时计算分解矩阵时会产生不可接受的巨大开销

提出方案：

• 线性投影的权重矩阵 W 分解为两个低秩矩阵A和B。输入 X 被下投影到潜在表示 H，并缓存。Y 可以通过上投影矩阵 B 从 H 重建。(也就是说，KV缓存变成了H，它是一个L*k的tensor，相比于之间的空间占用是k/d)

子问题：

• 所有注意力头一块分解可以提高准确性，但是因为要分解一个d*d的矩阵，导致分解和重建成本增加；每个头分别重建成本较低，但是准确率损失（可能是由于 SVD 无法捕捉到头之间共享的公共信息）。因此提出了中粒度的按组分解，像GQA那样

• 每个线性投影模块对压缩的敏感性不同，如果是敏感度高的矩阵需要更高的秩。因此设计了基于Fisher 信息的有效秩搜索算法。Fisher信息越高就越敏感，Fisher信息是Fisher得分函数的方差，Fisher得分是对数似然函数对参数​的一阶导。

Palu的整体流程：

一些优化：

• 矩阵吸收：W_v分解后的B_v 可以被W_o吸收，W_k分解后的A_k可以被W_q吸收。位置编码RoPE会阻止注意力部分矩阵吸收，为了解决这个问题，Palu 在解码过程中动态地从潜在表示中重建键。通过一个定制的 GPU 内核提高了重建效率

• 矩阵分解：整个MHA一起分解虽然能保证高准确率，但是会带来高计算量和高显存占用；因此使用了组矩阵分解，即若干个注意力头为一组，以组为单位进行矩阵分解

• 秩的选择：Fisher信息可以量化每个参数的信息量，假设压缩敏感性与 Fisher 信息成正比，通过计算每个权重矩阵的 Fisher 信息与所有分解目标的 Fisher 信息总和的比率来确定每个权重矩阵的秩。

• 处理SVD分解后异常值：SVD分解后会有异常值，这会导致难以量化。受量化文献的启发，应用瓦舍-哈达玛变换（WHT）来消除异常值。之后将哈达玛矩阵吸收到低秩分解权重A和B中，因此无需额外的计算开销就能优化此过程

实验：

• Palu方法和量化结合使用，KV缓存保留为原来的70%和50%，再加上使用量化方法，能够使得KV缓存的显存占用量为原本的13.215%

• Palu方法(-30%)和atom方法、KV Quant方法比较，在wikitext2的困惑度更低、显存占用更低、加速率更高

• 在长文本任务下，可KIVI-2比较，Palu 实现了相似的准确率(大部分情况下KIVI-2准确率更高)，同时具有额外的 30%低秩压缩率

• 在加速比方面，和KIVI-4bit比较，Palu-4bit的加速比高很多