牛叉！MIT 提出LongLoRA，微调LLaMA2，Token从4K扩展至100K

引言

 今年5月份，华盛顿大学提出了QLoRA，该算法在保持完整的16位微调任务性能的情况下，可以实现单卡48G GPU微调650亿参数的大模型，可谓是红极一时引爆了学术圈。

 今天这篇文章同样厉害，他们提出了「LongLoRA微调算法」，它能够在资源受限的情况下，极大的扩展预训练大模型（LLMs）的上下文长度，「LongLoRA可以让LLaMA2-7B的上下文从4K扩展至100K」，除此之外，作者公布了一个长文本 LongQA数据集，包含超过 3k 个长上下文问答对，可用于用于监督微调。

Paper：https://arxiv.org/pdf/2309.12307v1.pdf

Code：https://github.com/dvlab-research/longlora

背景介绍

 大型语言模型 (LLM) 通常使用固定的上下文Token进行训练，例如LLaMA采用2048个Token，LLaMA2 采用4096个Token。然而，上下文大小的限制了大模型在许多场景中的应用，比如长文档总结、长问题回答等。为了打破这一限制，最新的一些研究人员开始尝试通过对LLM进行训练微调来适配较长上下文场景。

 但通过长序列数据集从头开始训练一个大模型是非常具有挑战性的，并且对现有预训练的 LLM 进行微调也相当昂贵。例如，Position Interpolation使用了32 个A100 GPU将 LLaMA 模型从 2k 上下文扩展到 8k，并使用 128 个 A100 GPU 进行更长的上下文微调；FOT使用 32 个 TPU 进行标准Transformer 训练，使用 128 个 TPU 增加LLaMA上下文长度。这些资源对于绝大部分人来说是无法使用到的。为此作者考虑寻求有效的算法来扩展LLMs的上下文长度。

 一种直接有效的方法是通过低秩适应 (LoRA) 微调预训练的 LLM。LoRA 利用低秩矩阵修改自注意力块中的线性投影层，这通常很有效并且减少了可训练参数的数量。然而，实验结果表明，以这种方式训练的长上下文模型效果并不好且效率低。

「在效果方面」，简单的低秩自适应会导致长上下文扩展中的高度困惑，如下表所示将秩增加到更高的值并不能缓解此问题；

「在效率方面」，无论是否采用 LoRA，随着上下文大小的扩展，计算成本都会急剧增加，这主要是由于标准的自注意力机制，如下图所示，即使使用 LoRA，当上下文窗口扩大时，标准 LLaMA2 模型的训练时间也会大幅增加。

LongLoRA

 基于以上分析，本文作者提出了LongLoRA微调方法，它可以扩展预训练LLM的下文长度，例如：LLaMA、LLaMA2等。在一台 8× A100 机器上，微调后的LLaMA2-7B模型上下文长度可以达到100k，微调后的LLaMA2-70B模型上下文长度可以高达 32k 。

 LoRA的主要工作原理是通过使用低秩权重更新来近似完全微调，地作者发现短时间的注意力也能够在训练过程中近似长上下文。为此作者提出短时转移注意(

)来替代标准自注意力。如下图所示：

 将上下文长度分为几个组，并在每个组中分别进行注意力处理。在半注意力头中，将token移动组大小的一半，这确保了相邻组之间的信息流。例如，可以使用组大小为 2048 的

来近似训练总共 8192 个Token上下文长度。

微调之后的模型在推理过程中保留了原始的注意力架构。一些常见的LLM的技术也可以与该方法进行匹配。例如：FlashAttention-2在训练和推理时间上都与本文方法兼容，其主要原因是短注意力类似于LLM在预训练阶段的注意力方案。

 实验表明，可学习的嵌入和归一化层是解锁长上下文 LoRA 微调的关键，如下表所示。嵌入和归一化层在整个 LLM 中只占参数的一小部分。例如，LLaMA2 7B 中的嵌入具有 (< 2%) 个参数，归一化具有 (≤ 0.004%) 个参数。对于更大的预训练模型，这个比率会下降。

 除此之外，作者还提出了一个用于监督微调（SFT）的数据集 LongQA。LongQA 包含超过 3k 个长问题和相应的答案。它包含技术论文、科幻小说等书籍的各种类型的问题。例如摘要、关系、人物以及与材料相关的其它细节。

实验结果

 为了证明LongLoRA 的有效性和高效性。作者展示了扩展 LLaMA2 7B、13B 和 70B 上下文窗口的实验结果。根据位置插值的实验设置，使用适当的位置嵌入对模型进行微调。经过训练的模型实现了与全注意力和完全微调结果相当的性能，而计算成本要低得多，如图下图所示。

 利用RedPajama数据集进行训练，使用 LongLoRA 微调的模型可以发现随着评估上下文长度的增加复杂度逐渐降低。

 利用LongLoRA微调得到的LLaMA2模型，与其它开源模型进行对比。可以发现与最先进的模型LongChat-13B性能相当。