摘要：Transformers已成为大型语言模型（LLM）的支柱。然而，由于需要在内存中存储过去标记的键值表示缓存，其大小与输入序列长度和批量大小成线性比例，因此生成效率仍然很低。作为解决方案，我们提出了动态内存压缩（DMC），这是一种在推理时在线压缩键值缓存的方法。最重要的是，该模型可以学习在不同的头和层中应用不同的压缩率。我们将预训练的 LLM（如 Llama 2（7B、13B 和 70B））改装成 DMC Transformers，在英伟达 H100 GPU 上实现了高达 ~3.7 倍的自动回归推理吞吐量提升。DMC 通过持续预训练应用于可忽略不计的原始数据百分比，而不添加任何额外参数。我们发现，DMC 保持了原有的下游性能，缓存压缩率高达 4 倍，优于向上训练的分组查询注意（GQA）。GQA 和 DMC 甚至可以结合使用，以获得复合增益。因此，在任何给定的内存预算内，DMC 都能适应更长的上下文和更大的批次。

缓解自回归生成中的内存瓶颈

这篇论文的核心，是瞄准了大型语言模型（LLMs）在推理阶段一个非常实际的问题——内存消耗。简单来说，当模型一个字一个字地生成文本时，它必须在内存里存下所有过去token的键值对（Key-Value）表示。这个缓存的大小，和输入序列的长度以及批量大小呈线性关系。序列一长，或者处理的任务一多，内存就撑不住了，生成效率自然就会慢下来。

问题出在哪里？现有的解决方案，比如分组查询注意力（GQA），确实能通过共享参数来减少内存中的key-value头数量，但代价往往是牺牲了一些下游任务的性能。这是个典型的“既要马儿跑，又要马儿少吃草”的困境。

DMC：让模型自己决定怎么压缩

论文提出的动态内存压缩（DMC），思路很巧妙。它不是在推理前就预设好一个固定的压缩策略，而是让模型在生成过程中，动态地决定“是否要把当前的键值表示追加到缓存里”，还是“跟缓存里已有的顶端元素做一个加权平均”。

这背后的机制，具体是怎么实现的？

• 在线压缩：模型通过一个决策变量，对输入序列进行分段。每个段落可以独立决定是继续累积还是直接追加。这有点像把一长串信息，按内容的重要性和相关性，打包成几个精炼的“包裹”。
• 重要性评分：DMC还引入了一个重要性变量。在做加权平均时，模型会给当前token和缓存顶部元素分配不同的权重，确保更关键的信息被保留下来。
• 端到端学习：最有趣的是，这一切学习过程，都只通过在原始数据的一小部分上进行持续预训练来完成，且不增加任何新参数。模型会自己学会在不同层、不同头之间，应用差别化的压缩率。比如，某些层可能只需要1.5倍压缩，而另一些层则可以压缩到4倍。

实验结果：不仅省内存，性能还扛得住

论文在Llama 2的7B、13B和70B模型上做了大量实验。结果非常直观：

• 效率提升：在H100 GPU上，实现了高达3.7倍的推理吞吐量提升。这意味着相同时间内能处理更多请求，或者支撑更长的上下文。
• 性能保持：在4倍压缩率下，模型在MMLU、常识推理、代码生成等任务上的表现，基本和原始模型持平，甚至还有所提升。
• 与GQA的对比：与GQA相比，DMC在相似或更少的训练步骤下，达到了更好的性能。更重要的是，两者可以结合使用，产生“1+1>2”的复合增益。

未来探索方向：不止于Llama

DMC的潜力显然不止于此。以下几个方向值得关注：

• 跨架构适用性：这套方法能否迁移到BERT、T5等其他Transformer变体上？甚至是更广泛的序列处理模型？
• 预训练阶段集成：如果在模型预训练的早期就直接引入DMC，其效果和效率会不会有更根本性的提升？
• 多场景验证：在超长序列（比如整本书）、高并发实时生成（比如聊天机器人）等更复杂的实际场景中，它的表现如何？
• 与技术组合拳：将DMC与量化、剪枝、知识蒸馏等成熟的模型压缩技术结合，很可能会产生更惊人的效果。

总的来说，DMC并不是一个空中楼阁的理论构想，而是一个非常务实、且经得起检验的工程解决方案。它在不损害模型表现的前提下，切中要害地解决了推理阶段的内存效率问题。对于那些正在用大模型做长文本生成、或者需要服务大规模用户的应用来说，DMC提供了一条非常实用的技术路径。毕竟，在AI模型参数量不断攀升的今天，如何让它们“反赌、吃得少”，才是落地的关键所在。