arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2401.03462cs.CLcs.AI

使用 Activation Beacon 进行长上下文压缩

Long Context Compression with Activation Beacon

Beijing Academy of Artificial Intelligence(北京人工智能研究院);Gaoling School of Artificial Intelligence, Renmin University of China(中国人民大学高瓴人工智能学院)

Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, Zhicheng Dou

AI总结:

本文提出插件模块 Activation Beacon,通过直接压缩激活、定制工作流及压缩自回归训练,实现 LLMs 长上下文的高效压缩,在保持性能的同时显著降低内存开销并加速推理。

摘要:

长上下文压缩是一个关键的研究问题,因为它在降低与 LLMs 相关的高昂计算和内存成本方面具有重要意义。在本文中,我们提出了 Activation Beacon,这是一个用于基于 transformer 的 LLMs 的插件模块,旨在实现长上下文的有效、高效和灵活压缩。为了实现这一目标,我们的方法引入了以下技术设计。1) 我们直接压缩激活(即每一层的键和值),而不是利用软提示来传递信息(这构成了将长上下文中的复杂信息封装起来的主要瓶颈)。2) 我们定制了压缩工作流,其中每个细粒度的输入单元被逐步压缩,从而在训练和推理期间实现高质量压缩和高效计算。3) 我们通过基于压缩的自回归来训练模型,充分利用纯文本和指令数据来优化模型的压缩性能。4) 在训练期间,我们在每一步随机采样一个压缩率,教导模型支持广泛的压缩配置。我们在各种长上下文任务上进行了广泛的评估,这些任务的长度(例如 128K)可能远远超过最大训练长度(20K),例如文档理解、少样本学习和大海捞针。虽然现有方法难以处理这些具有挑战性的任务,但 Activation Beacon 在各种场景中保持了与未压缩基线相当的性能,在推理时间上实现了 2 倍的加速,并在 KV cache 的内存成本上实现了 8 倍的降低。我们的数据、模型和代码已在 \url{https://github.com/FlagOpen/FlagEmbedding/} 发布。

英文摘要:

Long context compression is a critical research problem due to its significance in reducing the high computational and memory costs associated with LLMs. In this paper, we propose Activation Beacon, a plug-in module for transformer-based LLMs that targets effective, efficient, and flexible compression of long contexts. To achieve this, our method introduces the following technical designs. 1) We directly compress the activations (i.e. keys and values at every layer), rather than leveraging soft prompts to relay information (which constitute a major bottleneck to encapsulate the complex information within long contexts). 2) We tailor the compression workflow, where each fine-grained input unit is progressively compressed, enabling high-quality compression and efficient computation during both training and inference. 3) We train the model through compression-based auto-regression, making full use of plain texts and instructional data to optimize the model's compression performance. 4) During training, we randomly sample a compression ratio at each step, teaching the model to support a wide range of compression configurations. Extensive evaluations are conducted on various long-context tasks whose lengths (e.g., 128K) may far exceed the maximum training length (20K), such as document understanding, few-shot learning, and Needle-in-a-Haystack. Whilst existing methods struggle to handle these challenging tasks, Activation Beacon maintains a comparable performance to the uncompressed baseline across various scenarios, achieving a 2x acceleration in inference time and an 8x reduction of memory costs for KV cache. Our data, model, and code have been released at \url{https://github.com/FlagOpen/FlagEmbedding/}.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑