A JoLT for the KV Cache: Near-Lossless KV Cache Compression via Joint Tucker and JL-Residual Allocation for LLMs
JoLT 通过对 Token 轴和特征轴应用部分 Tucker 分解，并利用经过 Johnson-Lindenstrauss 旋转的低比特残差来恢复丢弃的信息，在统一的字节预算下进行了优化，从而实现了大语言模型 KV Cache 近乎无损的 2-3 倍压缩，同时保持了困惑度（perplexity）和下游任务上的基准性能。
原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。 阅读完整免责声明
想象一下你正在运行一个庞大的、超级聪明的机器人大脑（大型语言模型），它正试图讲述一个故事。为了让故事持续进行，机器人必须记住它在当前对话中说过的每一个词。它将这些记忆保存在一个特殊的“笔记本”中，叫做 KV Cache。
问题在于，随着故事变得越来越长，这个笔记本变得越来越大。它开始吞噬掉所有的内存，让一切都慢了下来。这就像是在马拉松比赛中背着一座图书馆在奔跑；最终，你会停下来。
科学家们以前也尝试过缩小这个笔记本。有人尝试把整个东西挤压进一个小盒子（量化），也有人尝试通过寻找行或列中的模式来对页面进行总结（低秩方法）。但本文的作者 Rahul Krishnan 和 Volker Schulz 注意到，这些其他方法忽略了这一点：这个笔记本不仅仅是一个扁平的纸堆。它是一个 3D 的块体，拥有三个不同的面：头（Heads）（不同的思考方式）、标记（Tokens）（单词）和特征（Features）（细节）。
他们发现，这两个面充满了可以被轻易挤压的无聊、重复的内容，而另外两个面则是独特的，不能被挤压，否则会损失机器人的脑力。
重大发现：“JoLT”方法
团队发明了一个被称为 JoLT（联合 Tucker 和 JL-残差分配）的新技巧。把它想象成一个超级聪明的背包打包服务。
- 智能挤压（部分 Tucker）： JoLT 不尝试挤压整个 3D 块，而是观察数据并说：“好吧，‘头’和‘层’是独特且珍贵的；让我们不去动它们。但是‘标记’和‘特征’充满了废话。”因此，它只挤压这两个特定的面。这就像是拿一个巨大的、蓬松的枕头，只把中间的空气挤压出来，同时保持坚固的边缘完好无损。
- 安全网（JL-残差）： 当你挤压一个枕头时，一些空气会逃逸。如果你就那样不管，枕头就会变扁且毫无用处。JoLT 捕捉这些“逃逸的空气”（丢失的信息），并将其存储在一个微小的、超高效的安全网中，称为 JL-residual。这个网的打包能力非常出色，可以用极少的比特数保存缺失的细节。
- 完美的平衡（拉格朗日对偶）： 这是神奇之处。机器人有一个严格的空间预算（例如 1 字节）。JoNT 使用一个数学上的“智能分配器”来决定究竟要挤压多少枕头，以及要给安全网分配多少空间。它意识到某些部分的记忆（“键/Keys”）很容易挤压，而另一些部分（“值/Values”）很顽固，需要更多的安全网空间。它动态地移动预算，以获得最佳结果。
他们证明了什么（以及没能证明什么）
作者在两个著名的机器人大脑上测试了这一点：Mistral-7B（使用“分组查询”风格）和 LLaMA-2-13B（使用“多头”风格）。
“免费区”： 他们发现了一个甜点区，在这里他们可以将内存缩小 2 到 3 倍（2–3×），且机器人的性能完全没有下降。这是“近乎无损的”。
- 在 GSM8K（数学问题）和 RULER（在文本大海中寻找针头）等测试中，压缩后的机器人得分与未压缩的机器人完全相同，处于统计噪声范围内。
- 重建记忆的误差极小：键（Keys）约为 0.009，值（Values）约为 0.006。这比之前的各种方法（如 4-bit 量化或跨层 SVD）大约好 10 倍（一个数量级）。
“悬崖”： 他们还发现了一个极限。如果你试图过度挤压内存（超过 3 倍），情况就会变得混乱。
- Mistral 机器人表现出优雅的退化，随着挤压程度增加，性能缓慢下降。
- 然而，LLaMA 机器人却在 4 倍到 5 倍压缩之间撞上了“悬崖”。它的性能大幅崩溃，从 5.39 的得分骤降至 9.07（困惑度发生了巨大跳跃，意味着它预测单词的能力变差了很多）。
快速版本：FlashJoLT
计算完美的挤压需要时间。为了解决这个问题，他们创建了 FlashJoLT。与其每次都进行沉重的完美数学计算，它使用一种“随机化”的捷径来快速猜测主要模式。
- 结果： 它在压缩内存方面快了 5 到 13 倍，但质量与缓慢的完美版本完全保持一致。
他们排除了哪些可能
论文非常明确地说明了哪些方法在处理这个特定问题时效果不佳：
- 挤压一切： 尝试压缩所有三个面（头、标记和特征）是个坏主意。“头”和“层”太独特了；挤压它们会损害机器人的大脑。
- 固定比特量化： 仅仅降低每个数字的比特数（比如强制全部为 4 比特）无法达到 2–3× 压缩的“甜点区”。它要么压缩得不够，要么损失了太多质量。
- 一刀切： 你不能把“键（Keys）”和“值（Values）”同等对待。“值”要难压缩得多（难 2–3 倍），因此它们需要不同的空间预算。
底线
作者在真实硬件（A100 GPU）上测量了这一点，发现 JoNT 提供了一种近乎无损的方式，可以在不损害机器人智能的情况下将内存缩小 2–3 倍。
然而，他们谨慎地指出，这并不是解决所有问题的万灵药。
- 它在“免费区”（2–3×）表现出色，但如果对某些类型的机器人（如 LLaMA）施加更强的压力，会导致质量剧烈下降。
- 虽然内存存储变小了，但机器人每次说话时仍需进行一些数学运算来“解压”记忆。他们建议，为了使这在现实世界中真正实用，工程师需要构建特殊的计算机芯片（融合算子/fused kernels），以便能够直接读取挤压后的内存，而无需先进行解压。
简而言之，JoLT 是一个精妙的、基于数学的打包技巧，它为长对话节省了大量的空间，但它存在极限，并且需要未来硬件的辅助才能实现全速运行。
您所在领域的论文太多了？
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要，使用您的语言。