LoCoCo: Dropping In Convolutions for Long Context Compression

📄 arXiv: 2406.05317v2 📥 PDF

作者: Ruisi Cai, Yuandong Tian, Zhangyang Wang, Beidi Chen

分类: cs.LG, cs.CL

发布日期: 2024-06-08 (更新: 2024-10-25)

💡 一句话要点

提出LoCoCo，通过卷积动态融合KV缓存，实现长文本高效压缩。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本压缩 键值缓存 卷积神经网络 大型语言模型 上下文建模

📋 核心要点

  1. 现有长文本处理方法在LLM中面临内存瓶颈，选择性丢弃KV对可能导致上下文信息丢失。
  2. LoCoCo通过卷积融合新旧KV对，自适应地保留重要上下文信息，实现高效压缩。
  3. 实验表明，LoCoCo在推理和微调阶段均能有效压缩长文本，并保持或提升模型性能。

📝 摘要（中文）

本文提出了一种名为LoCoCo（Dropping In Convolutions for Long Context Compression）的新方法，旨在解决大型语言模型（LLM）处理长文本序列时面临的内存瓶颈。LoCoCo仅使用固定大小的键值（KV）缓存，从而提高推理和微调阶段的效率。与之前基于启发式方法选择性丢弃KV对的方法不同，LoCoCo采用数据驱动的自适应融合技术，将先前的KV对与传入的token融合，以最大限度地减少上下文信息的丢失，并确保准确的注意力建模。这种token集成通过注入一维卷积核来实现，该卷积核动态计算每个KV缓存槽的混合权重。LoCoCo设计与现有LLM框架广泛兼容，无需架构修改即可轻松“嵌入”，同时产生最小的调整开销。实验表明，LoCoCo在各种上下文长度下均保持出色的性能，并且在推理和微调阶段均可实现高上下文压缩率。在推理过程中，我们成功地将多达3482个token压缩到128大小的KV缓存中，同时保持与完整序列相当的性能，与相同缓存大小的基线相比，准确率提高了0.2791。在后训练调整期间，我们还使用固定大小为512的KV缓存有效地将上下文长度从4K扩展到32K，实现了与使用完整序列进行微调相似的性能。

🔬 方法详解

问题定义：大型语言模型处理长文本时，KV缓存占用大量内存，限制了模型能够处理的上下文长度。现有方法通过启发式规则丢弃部分KV对，但可能导致关键上下文信息丢失，影响模型性能。

核心思路：LoCoCo的核心思路是利用卷积操作，将新输入的token信息与已有的KV缓存进行融合，从而在固定大小的KV缓存中保留尽可能多的上下文信息。通过数据驱动的方式学习融合权重，自适应地决定哪些历史信息应该保留，哪些应该被新信息覆盖。

技术框架：LoCoCo的主要流程如下：1. 输入长文本序列；2. 使用固定大小的KV缓存存储上下文信息；3. 对于每个新输入的token，使用一维卷积核计算其与KV缓存中每个槽的融合权重；4. 根据计算出的权重，将新token的信息融合到KV缓存中，更新缓存内容；5. 使用更新后的KV缓存进行后续的注意力计算和模型推理。

关键创新：LoCoCo的关键创新在于使用卷积操作进行KV缓存的融合。与传统的KV缓存更新方式不同，LoCoCo不是简单地替换或丢弃KV对，而是通过卷积核学习到的权重，将新token的信息与已有信息进行加权融合，从而更有效地保留上下文信息。这种数据驱动的融合方式能够自适应地选择哪些信息应该保留，哪些应该被更新，从而在有限的缓存空间内最大化信息的利用率。

关键设计：LoCoCo的关键设计包括：1. 使用一维卷积核进行融合，降低计算复杂度；2. 卷积核的权重通过训练数据学习得到，实现自适应的融合；3. KV缓存的大小固定，保证了内存占用可控；4. LoCoCo可以作为“drop-in”模块集成到现有的LLM框架中，无需修改模型架构。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

LoCoCo在推理阶段成功将3482个token压缩到128大小的KV缓存中，性能与完整序列相当，相比基线方法在相同缓存大小下准确率提升高达0.2791。在后训练调整阶段，使用固定大小为512的KV缓存，将上下文长度从4K扩展到32K，实现了与使用完整序列进行微调相似的性能。这些结果表明LoCoCo能够有效压缩长文本，并在保持甚至提升模型性能的同时，显著降低内存需求。

🎯 应用场景

LoCoCo可应用于各种需要处理长文本的场景，例如：长文档摘要、机器翻译、对话系统、代码生成等。通过压缩上下文信息，LoCoCo可以降低模型的内存需求，提高推理速度，并支持处理更长的输入序列。该技术有助于在资源受限的设备上部署大型语言模型，并提升长文本处理任务的性能。

📄 摘要（原文）

This paper tackles the memory hurdle of processing long context sequences in Large Language Models (LLMs), by presenting a novel approach, Dropping In Convolutions for Long Context Compression (LoCoCo). LoCoCo employs only a fixed-size Key-Value (KV) cache, and can enhance efficiency in both inference and fine-tuning stages. Diverging from prior methods that selectively drop KV pairs based on heuristics, LoCoCo leverages a data-driven adaptive fusion technique, blending previous KV pairs with incoming tokens to minimize the loss of contextual information and ensure accurate attention modeling. This token integration is achieved through injecting one-dimensional convolutional kernels that dynamically calculate mixing weights for each KV cache slot. Designed for broad compatibility with existing LLM frameworks, LoCoCo allows for straightforward "drop-in" integration without needing architectural modifications, while incurring minimal tuning overhead. Experiments demonstrate that LoCoCo maintains consistently outstanding performance across various context lengths and can achieve a high context compression rate during both inference and fine-tuning phases. During inference, we successfully compressed up to 3482 tokens into a 128-size KV cache, while retaining comparable performance to the full sequence - an accuracy improvement of up to 0.2791 compared to baselines at the same cache size. During post-training tuning, we also effectively extended the context length from 4K to 32K using a KV cache of fixed size 512, achieving performance similar to fine-tuning with entire sequences.