LONGWRITER：从长篇上下文 LLM 中释放 10,000 多单词的生成能力

24年8月来自清华和智谱AI公司的论文“LONGWRITER: UNLEASHING 10,000+ WORD GENERATION FROM LONG CONTEXT LLMS”。

当前的长上下文大语言模型 (LLM) 可以处理多达 100,000 个tokens的输入，但生成甚至超过 2,000 个单词长度的输出却举步维艰。受控实验发现，模型的有效生成长度本质上受其在监督微调 (SFT) 期间看到的样本限制。换句话说，它们的输出限制是由于现有 SFT 数据集中长输出示例的稀缺性。为了解决这个问题，引入 AgentWrite，这是一个基于智体的流水线，它将超长生成任务分解为子任务，使现成的 LLM 能够生成超过 20,000 个单词的连贯输出。利用 AgentWrite，构建 LongWriter-6k，这是一个包含 6,000 个 SFT 数据的数据集，输出长度从 2k 到 32k 单词不等。将此数据集纳入模型训练，现有模型的输出长度扩展到 10,000 字以上，同时保持了输出质量。还开发 LongBench-Write，这是一个用于评估超长生成能力的综合基准。原 9B 参数模型通过 DPO 进一步改进，在此基准上实现了最先进的性能，甚至超越了更大的专有模型。

代码和模型位于：https://github.com/THUDM/LongWriter。

长上下文大语言模型 (LLM) 的最新进展推动了内存容量显著扩展的模型开发，这些模型能够处理长度超过 100,000 个tokens的历史记录（Anthropic，2024；Reid，2024；GLM，2024）。然而，尽管它们能够处理大量输入，但当前的长上下文 LLM 难以生成同样长的输出。

构建 LongWrite-Ruler 评估来探究 LLM 的生成长度限制。探索其生成长度受限的原因：改变模型 SFT 阶段数据的最大输出长度，发现在 LongWrite-Ruler 测试中训练的模型的最大输出长度与 SFT 数据的最大输出长度呈现出显著的正相关性。

为了探究 LLM 可以提供的最大输出长度，构建了一个轻量级测试：创建了 8 个不同的指令，中文和英文各 4 个，并在指令中改变输出长度要求“L”。例如，“写一篇关于罗马帝国历史的 L 字文章”。在测试期间，使用 L ∈ {1000, 2000, 5000, 10000, 20000, 30000}，总共得到 48 个测试提示。

在 LongWrite-Ruler 上测量了 4 个开源模型和 4 个专有模型的最大输出长度。在推理期间，将温度设置为 0.5。对于专有模型，将生成的 max tokens 参数配置为相应模型 API 调用支持的最大输出长度。对于开源模型，将其设置为 32k。在输出中，验证没有模型由于 max tokens 约束而产生截断的输出，这可能低估了它们的最大输出长度。同时，几乎没有观察到重复内容生成的情况，这可能导致高估。结果如图所示：对于每个长度要求（x 轴），绘制模型在 8 条相应指令上的平均输出长度（y 轴）。对 x 轴和 y 轴使用对数刻度。从图中可以看出，所有模型的最大输出长度约为 2k 个字。专有模型的有效输出窗口大小通常无法达到其最大 token 生成长度。此外，由于拒绝情况的增加，随着所需长度超过 10k，平均输出长度甚至会减少。

从LongWriter-Ruler的测试结果来看，所有模型都无法产生长度超过 2,000 字的输出。同时，对 WildChat 用户交互日志的分析（Zhao，2024）表明，超过 1% 的用户提示明确要求输出超出此限制，这凸显了当前研究迫切需要克服这一限制。

假设常见的 2,000 字输出长度限制是由于 SFT 数据中固有的输出长度限制，即“一个人只有读过书才能说话”。为了检验这一假设，改变 SFT 数据进行一系列受控实验。在实验中，用 GLM-4-9B（GLM，2024）作为基础模型，并选择 GLM-4 的聊天 SFT 数据（总共 180k 个数据，是 GLM-4 整个 SFT 数据的子集）作为完整的 SFT 数据集。为了控制 SFT 数据的最大输出长度，分别过滤掉输出长度超过 500、1,000 和 2,000 字的数据。这产生了三个训练集，分别包含原始数据的 72%、98% 和 99.9%。

在这三个训练集上训练 GLM-4-9B 模型，并在 LongWriter-Ruler 上测量所得模型的最大输出长度（使用 L ∈ {500, 1000, 2000, 4000} 进行测试）。如图所示，模型的最大输出长度与 SFT 数据中的最大输出长度成比例增加，分别达到约 600、900 和 1,800 个字。最大输出长度的增加也对应于模型在每个所需长度指令的平均输出长度改善。这一发现表明，模型的输出限制是由于 SFT 数据中的输出长度不足造成的。此外，这个限制无法通过 LLM 合成训练数据（Tunstall，2023；Abdin，2024）或迭代 SFT（Chen，2024b；Burns，2023）克服，因为现有模型生成的数据仍然无法突破长度限制。

为了利用现成的 LLM 自动生成具有较长输出的 SFT 数据，设计了 AgentWrite，这是一种分而治之式的智体流水线（如图所示）。AgentWrite 首先将长篇写作任务分解为多个子任务，每个子任务仅要求模型写一段话。然后，模型按顺序执行这些子任务，将子任务输出连接起来以获得最终的长篇输出。这种使用 LLM 智体将复杂任务分解为多个子任务的方法已经应用于各个领域，例如解决问题（Wu et al., 2023）、软件开发（Qian et al., 2023）和模型评估（Saha et al., 2024）。该工作探索整合规划以使模型能够完成复杂的长篇写作任务工作。

第一步：规划
灵感来自人类作家的思维过程，作家通常首先制定长篇写作任务的总体计划，通常涉及概述结构并规划每个部分的内容和长度。利用 LLM 的规划功能，根据写作说明输出这样的写作大纲，其中包括每个段落的主要内容和字数要求。

第二步：写作
从步骤 I 获得写作计划后，串行调用 LLM 完成每个子任务，逐段生成写作内容。为了保证输出的连贯性，在调用模型生成第 n 段时，还输入了之前生成的 n-1 段，让模型根据已有的写作历史继续写作下一段。虽然这种串行方式无法并行调用模型同时完成多个子任务，输入长度也会变长，但在验证中表明，以这种方式获得的写作整体连贯性和质量远远优于并行生成的输出。

第三步：验证
在两个长篇写作数据集上测试了 AgentWrite 方法的生成长度和质量。第一个数据集是 LongWrite-Ruler，用于准确测量该方法可以提供的输出长度。第二个数据集是构建的 LongBench-Write 基准，主要用于评估模型生成的内容在长度和写作质量方面与用户指示的匹配程度。

有了一个利用现成 LLM 自动生成更长输出的智体框架，问题是：是否可以将这种生成超长输出的能力教给 LLM，让它们在一次输出中完成长篇写作任务？带着这个问题，进行模型训练实验。

基于两个最新的开源模型 GLM-4-9B1 和 Llama-3.1-8B2 进行训练。这两个模型都是基础模型，支持高达 128k 个 token 的上下文窗口，因此非常适合在长输出上进行训练。为了提高训练效率，采用了带损失加权的packing训练 (Bai et al., 2024a)。对这两个模型的训练产生了两个模型：LongWriter-9B（GLM-4-9B-LongWriter 缩写）和 LongWriter-8B（Llama-3.1-8B-LongWriter 缩写）。

如果按序列平均损失，即取批次内每个序列平均损失的平均值，则每个目标 token 对长输出数据损失的贡献将明显小于那些输出较短的数据。在实验中，这会导致模型在输出较长的任务上性能不佳。因此，选择了一种损失加权策略，即按 token 平均损失，其中损失计算为该批次内所有目标 token 的损失平均值。
所有模型均使用具有 8xH800 80G GPU 和 DeepSpeed+ZeRO3+CPU 卸载的节点进行训练（Rasley，2020 年）。使用的批次大小为 8，学习率为 1e-5，包装长度为 32k。对模型进行 4 个 epoch 的训练，大约需要 2,500-3,000 步。

为了进一步提高模型的输出质量并增强其遵循指令长度约束的能力，对监督微调的 LongWriter-9B 模型执行DPO直接偏好优化（Rafailov，2024）。 DPO 数据来自 GLM-4 的聊天 DPO 数据（约 50k 条）。此外，构建了 4k 对专门针对长格式写作说明的数据。具体来说，对于每条写作说明，从 LongWriter-9B 中抽取 4 个输出，并按照 Hou et al. (2024) 中的方法对这些输出进行评分。还结合了长度跟随分数。然后，选择得分最高的输出作为正样本，并从其余三个输出中随机选择一个作为负样本。生成的模型 LongWriter-9B-DPO 在上述数据混合上进行了 250 步训练。按照 Hou et al. (2024) 中的方法进行 DPO 训练。