UIO-LLMs：面向长文本的大语言模型无偏增量优化方法UIO-LLMs：面向长文本的大语言模型无偏增量优化方法
摘要 大型语言模型（LLMs）在处理长文本时，由于上下文窗口大小的限制，面临着巨大挑战。本文介绍了一种名为UIO-LLMs的新方法，它是一种在长文本环境下对内存增强型Transformer进行无偏增量优化的方案。我们将这一过程概念化为一个简化的编码器-解码器框架，其中权重共享的编码器和解码器分别将上下文片段封装到内存中，并利用这些内存来预测后续片段的输出。随后，通过将内存增强型Transformer视为全连接递归神经网络（RNNs），我们使用截断时间反向传播（TBPTT）算法对训练过程进行了改进，该算法结合了创新的增量优化技术。这些技术不仅减少了时间复杂度，而且通过无偏优化过程解决了梯度计算中的偏差问题。UIO-LLMs成功地处理了长文本，例如将Llama2-7b-chat的上下文窗口从4K扩展到100K个token，而仅增加了2%的参数，同时随着上下文长度的增加，推理成本几乎呈线性增长。 关键词 上下文压缩 · 长文本LLMs 1. 引言 人们对大型语言模型（LLMs）[1, 2, 3]的长文本推理能力越来越感兴趣。LLMs的上下文窗口可以比作计算机的内存，更大的容量为开发者提供了更大的灵活性和可能性。这使得他们能够集成诸如检索增强生成（RAG）[4]等技术，并创建各种下游应用程序，如问答和阅读理解[5]。 然而，有限的计算资源使得在长文本上预训练模型几乎不可行。目前流行的方法是，首先使用短文本对模型进行预训练，然后通过微调扩展其处理长文本的能力。LongChat [6]、LongLora [7]、Positional Interpolation [8]、PoSE [9]、Yarn [10]等模型都采用了这种方法。然而，注意力机制固有的二次复杂度仍然是处理长文本时推理阶段效率的挑战。除了这些基于微调的方法外，另一种策略是在推理阶段进行适当的修改，以增加模型的有效上下文窗口大小。这些策略通常涉及注意力剪枝，例如Streaming LLM [11]，它通过只保留最近的KV缓存和最前面的KV缓存来管理token的数量。然而，对于这些基于剪枝的方法，来自丢弃token的信息变得难以利用，导致性能下降程度不同。 在本文中，我们研究并认识到，Transformer模型[12]通常会保留由于注意力机制而产生的完整历史信息集；相反，递归神经网络（RNNs）的特点是保留了提炼的历史信息，这是它们对序列数据处理的结果，强调了决策过程中最近的信息。在这方面，这两种架构表现出对比鲜明的特征。 某些技术，如Performer [13]和Linear Transformers [14]，通过采用核方法[15, 16]来修改注意力计算顺序。它们计算键和值的外部积，并将它们累加到一个大矩阵中进行数据压缩。这将Transformer转换为一个类似RNN的模型，该模型压缩所有过去的信息，削弱了其处理长期依赖关系的能力。在存储全面（Transformer）和压缩（RNN）历史数据之间取得平衡是可能的。 在这项研究中，我们提出了UIO-LLMs方法，如图1所示，该方法利用仅解码器LLMs作为上下文压缩器。具体来说，将上下文划分为多个片段，每个片段的末尾都附加了多个“”token。在编码器进行前向传播后，“”token的激活提炼了上下文信息，有效地形成了一个紧凑且信息丰富的内存表示。这种表示可以通过由两个投影矩阵组成的传输头作为额外的KV缓存传输到解码器。为了最大限度地减少引入额外的参数，我们利用LoRA [17]对编码器和传输头进行微调。这导致Llama2-7b-chat [18]的参数仅增加了2%。 关于优化，内存段的互连形成了类似于全连接RNN的结构。因此，时间反向传播（BPTT）对于优化至关重要。然而，它会导致线性时间和存储开销，并随着输入文本长度的增加而增加。因此，我们的研究重点是提高BPTT算法的效率。为此，我们引入了一种增量TBPTT算法，它是对截断BPTT方法[19]的改进，通过以增量方式重新排序计算过程，显著减少了时间开销。此外，尽管增量TBPTT提高了效率，但与局部TBPTT窗口相关的固有偏差梯度估计问题仍然是学习长期依赖关系的障碍。为了克服这一挑战，我们进一步开发了无偏增量优化算法。该算法确保了无偏梯度估计，促进了对长度高达100K的文本的训练，并具有恒定的压缩比。 [...]