LongAlign：大模型长文本处理能力提升之道

人工智能咨询培训老师叶梓 转载标明出处

大模型（LLMs）在处理长文本时，需要在输入序列上进行指令微调（instruction finetuning），以确保它们能够有效地处理长文本。现有的方法主要集中在上下文扩展（context extension），即通过位置编码扩展和对长文本的持续训练来提升模型的上下文长度。最近，一篇由清华大学和Zhipu.AI的研究人员共同撰写的论文《LongAlign: A Recipe for Long Context Alignment of Large Language Models》提出了一种新的方法，通过指令微调来提升LLMs处理长用户提示的能力。从而提高LLMs在长文本处理方面的能力。

在LongBench-Chat基准测试中的结果，包含了真实世界查询，长度在10k-100k之间

LongAlign

大模型（LLMs）通过监督式微调（Supervised Fine-Tuning, SFT）来学习对齐的基本概念程中，指令和响应被串联起来形成序列 [x, y]，然后输入到自回归语言模型中，目标是最大化给定指令 x 下产生响应 y 的概率 。损失函数的设计与传统语言模型相似，主要关注目标标记 y 的损失。

为了创建一个多样化的长指令数据集，研究者从九个不同的来源收集了长文本，包括书籍、百科全书、学术论文和代码等。利用Claude 2.1工具，根据这些长文本上下文生成任务和答案。为了增加生成任务的多样性，研究者在提示中加入了不同类型的任务描述，比如摘要、信息提取和推理等。通过这种方法，他们为10,000个长文本创建了任务和答案，形成了10,000个监督数据实例，其中10%是中文数据。这些数据的长度范围从8k到64k不等，由ChatGLM3-6B分词器进行测量，因为该分词器对中文字符有更高的压缩率。

数据构建的例子，说明了如何从长文本中生成任务和答案

在数据构建过程中，研究者特别注重数据多样性的重要性。他们通过从不同的文本来源收集数据，并使用不同的任务类型提示来生成问题，确保了数据集能够覆盖广泛的任务类型和语言。

为了确保模型在监督式微调后能够同时处理长文本和短文本（一般能力），研究者将长指令数据与一般指令数据集混合进行训练。在这种混合数据中，大量的一般短数据与相对较少的长指令数据混合，导致数据长度分布呈现长尾特征。为了最小化每个批次中的空闲时间，研究者探索了两种训练方法：打包（Packing）和排序批处理（Sorted Batching）。

Packing（打包）：涉及将不同长度的数据串联在一起，直到达到最大长度。然后将这些打包后的数据在多GPU上进行批处理和处理。为了防止同一打包中的不同序列在自注意力计算中发生交叉污染，研究者使用了 FlashAttention 2 中的 flash_attn_varlen_func 函数，该函数支持块对角注意力的高效计算。

在长尾数据长度分布下，通过打包或排序批处理可以减少空闲时间并加快训练过程

打包方法存在一个问题，即它倾向于更多地关注更长的序列和包含更多目标标记的序列。这是因为在最终损失计算中，不同打包包含的序列数量和目标标记数量不同，但它们对损失的贡献却被平等对待。为了解决这个问题，研究者们提出了一种损失加权策略，通过调整每个序列的损失权重，确保在损失计算中各个序列的贡献是平衡的。

Sorted Batching（排序批处理）：排序批处理方法则是另一种提高训练效率的策略。它通过将长度相似的数据序列放在同一批次中来进行训练，从而减少了批次内部的空闲时间。尽管这种方法可能会导致不同批次之间的数据分布偏差，但研究者们发现，通过使用大的梯度累积步长和优化器的强大适应性，这种方法在加速训练的同时并没有对性能产生明显的负面影响。

评估基准是LongBench-Chat ，包括50个长文本实际查询，长度从10k到100k不等，涵盖了文档问答、摘要和编码等多种关键用户密集型场景。这个基准包括40个英文任务和10个中文任务。为了确保评估真正反映了模型遵循长文本指令的能力，研究者避免了使用在预训练期间可能已经被模型看到并记忆的流行长文本，并避免了提出模型在不阅读长文本的情况下就能回答的问题。

使用 GPT-4 作为评估者，根据给定的人工注释参考答案和每个问题的少量评分示例，对模型的响应进行1-10的评分。为了确保评估者能够仅根据地面真实和少量评分示例做出明智的判断，研究者避免了过于开放式的问题，例如“根据前面的文本写一首诗”。为了验证使用 GPT-4 作为 LongBench-Chat 评估者的可靠性，研究者进行了人类评估研究，并在表格中展示了 GPT-4 的评估与人类评估者之间的相关性。

人类标注者之间的相关性，以及GPT-4（有无少量示例）与人类标注的相关性

实验

数据：研究者使用ShareGPT作为短指令数据的来源，并过滤掉空助手响应。为了比较长指令数据对模型训练的影响，他们包括了四种长指令数据集：'LongAlign-0k', 'LongAlign-5k', 'LongAlign-10k' 和 'LongAlpaca-12k'。这些数据集根据第3.2节中描述的过程构建，并且 'LongAlpaca-12k' 数据集提供了与 'LongAlign' 数据集相比较的多样性影响。

模型：研究者选择了三种模型变体：ChatGLM3-6B, Llama-2-7B 和 Llama-2-13B。考虑到这些基模型的上下文窗口大小，他们首先执行上下文扩展，将上下文窗口扩展到64k。

训练：所有模型使用8xA800 80G GPU和DeepSpeed+ZeRO3+CPU卸载进行训练。设置训练数据的最大长度为64k，超出此长度的数据将从右侧截断。

评估：评估包括长文本任务和短文本任务。使用了LongBench-Chat和LongBench评估模型的长文本对齐熟练度和一般长文本理解能力。对于短文本任务，使用了MT-Bench和Open LLM Leaderboard上的一般任务进行评估。

研究者对ChatGLM3-6B-64k进行了使用ShareGPT数据混合不同长指令数据集的监督式微调（SFT）。除了 'LongAlign-0k' 外，所有模型都使用更高效的打包策略和损失加权训练。评估结果显示，长指令数据的数量增加可以一致地提高模型在所有长任务中的性能，同时在短任务中的性能保持与仅使用短指令训练时相当的水平。

ChatGLM3-6B-64k模型在不同数量和类型的长指令数据训练后的性能

ChatGLM3-6B-64k模型在不同长数据集混合ShareGPT训练后，在1k-60k长度范围内的“针堆”测试性能

研究者比较了不同的训练方法，包括朴素批处理、打包（带和不带损失加权）、排序批处理，对ChatGLM3-6B-64k和Llama-2-6B-64k的训练效率和最终性能的影响。打包和排序批处理的训练效率是朴素批处理的两倍，并且在长任务和短任务上的表现与朴素批处理相当。损失加权策略显著提高了打包训练在长指令任务上的性能。

ChatGLM3-6B-64k和Llama-2-7B-64k模型在不同训练方法下的性能

在8xA800 80G GPU下，不同训练方法所需的训练时间

研究者探索了LongAlign框架在两个方向上的扩展性：更大的模型规模和更长的上下文窗口。他们使用LongAlign-10k数据集微调了Llama-2-13B-64k，并使用两种高效的训练方法，发现与7B规模的模型相比，13B模型在LongBench-Chat上的表现提高了10%。

在LLama-2-13B模型上进行LongAlign扩展的结果

研究者还展示了在模型训练过程中，长任务和短任务相对性能曲线的变化，说明了随着训练步骤的增加，性能如何变化。如图6所示长任务和短任务的学习曲线趋势非常相似，都显示出在0-500步之间快速提升，之后提升变慢，并在1000步后趋于稳定。

ChatGLM3-6B-64k模型在训练过程中长任务和短任务的相对性能变化

通过这些实验得出结论，LongAlign方法在数据量、数据多样性和正确的训练方法方面对最终性能至关重要。并且长任务和短任务的学习过程可能由共享的潜在因素共同决定，这些因素在训练过程中被优化，以帮助模型同时对齐长指令和短指令。

论文链接：https://arxiv.org/abs/2401.18058

项目链接：GitHub - THUDM/LongAlign: LongAlign: A Recipe for Long Context Alignment Encompassing Data, Training, and Evaluation