Soaring from 4K to 400K: Extending LLM's Context with Activation Beacon

LLM MQIA Long Context
2024年01月07日
长上下文的利用对于大型语言模型来说是一个巨大的挑战,因为它们的上下文窗口长度受到限制。虽然可以通过微调来扩展上下文窗口,但这将导致训练和推理时间的显著成本,并对LLM的原始能力产生不利影响。在本文中,我们提出了Activation Beacon,它将LLM的原始激活压缩成更紧凑的形式,从而使其能够使用有限的上下文窗口感知更长的上下文。Activation Beacon被引入作为LLM的即插即用模块。它在完全保留LLM在短上下文上的原始能力的同时,扩展了在处理长上下文方面的新能力。此外,它使用短滑动窗口来处理长上下文,从而在训练和推理中实现了竞争性的内存和时间效率。Activation Beacon是通过基于具有多样化压缩比的信标混合条件的自回归任务来学习的。由于这种处理方式,它可以仅使用短序列数据在短短的10K步内高效地训练,这在单个8xA800 GPU机器上不到9小时。实验研究表明,Activation Beacon能够将Llama-2-7B的上下文长度扩展100倍(从4K到400K),同时在长上下文生成和理解任务上实现了优异的结果。我们的模型和代码将在BGE代码库中提供。
The utilization of long contexts poses a big challenge for large language models due to their limited context window length. Although the context window can be extended through fine-tuning, it will result in a considerable cost at both training and inference time, and exert an unfavorable impact to the LLM's original capabilities. In this work, we propose Activation Beacon, which condenses LLM's raw activations into more compact forms such that it can perceive a much longer context with a limited context window. Activation Beacon is introduced as a plug-and-play module for the LLM. It fully preserves the LLM's original capability on short contexts while extending the new capability on processing longer contexts. Besides, it works with short sliding windows to process the long context, which achieves a competitive memory and time efficiency in both training and inference. Activation Beacon is learned by the auto-regression task conditioned on a mixture of beacons with diversified condensing ratios. Thanks to such a treatment, it can be efficiently trained purely with short-sequence data in just 10K steps, which consumes less than 9 hours on a single 8xA800 GPU machine. The experimental studies show that Activation Beacon is able to extend Llama-2-7B's context length by $\times100$ times (from 4K to 400K), meanwhile achieving a superior result on both long-context generation and understanding tasks. Our model and code will be available at the BGE repository.
许愿