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的加速比高很多