[论文解读] Extending Llama-3's Context Ten-Fold Overnight

Peitian Zhang, Ninglu Shao+5 more|arXiv (Cornell University)|Apr 30, 2024

Advanced Vision and Imaging被引用 4

一句话总结

该论文提出 Llama-3-8B-Instruct-80K-QLoRA，通过仅使用 GPT-4 生成的 3.5K 个合成训练样本和 QLoRA 微调，将 Llama-3 的上下文长度从 8K 扩展至 80K。该模型在长上下文基准测试中达到最先进性能，同时保持了短上下文任务的性能，仅在单台 8xA800（80GB）GPU 机器上训练 8 小时即可完成。

ABSTRACT

We extend the context length of Llama-3-8B-Instruct from 8K to 80K via QLoRA fine-tuning. The entire training cycle is super efficient, which takes 8 hours on one 8xA800 (80G) GPU machine. The resulted model exhibits superior performances across a broad range of evaluation tasks, such as NIHS, topic retrieval, and long-context language understanding; meanwhile, it also well preserves the original capability over short contexts. The dramatic context extension is mainly attributed to merely 3.5K synthetic training samples generated by GPT-4 , which indicates the LLMs' inherent (yet largely underestimated) potential to extend its original context length. In fact, the context length could be extended far beyond 80K with more computation resources. Therefore, the team will publicly release the entire resources (including data, model, data generation pipeline, training code) so as to facilitate the future research from the community: \url{https://github.com/FlagOpen/FlagEmbedding}.

研究动机与目标

• 实现 Llama-3-8B-Instruct 上下文长度从 8K 到 80K 的高效、低资源扩展，且计算开销极小。
• 探究大型语言模型是否能通过极少的合成数据，天然泛化至远超其原始上下文长度的场景。
• 在实现卓越长上下文能力的同时，保持原始模型在短上下文任务上的性能表现。
• 发布完整可复现的训练方案，包括数据、代码和模型权重，以加速社区在长上下文大语言模型领域的研究。
• 探索在增加计算资源的前提下，使用相同方法将上下文长度进一步外推至 80K 以上是否可行。

提出的方法

• 使用 GPT-4 合成 3.5K 个长上下文训练样本，涵盖三类任务：单细节问答（Single-Detail QA）、多细节问答（在同质与异质上下文上）以及传记摘要（Biography Summarization）。
• 通过组合来自同一长上下文（64K–80K token）的问答对，构建多轮对话，以模拟指令微调场景。
• 将合成数据与 5K 个 RedPajama 和 12K 个 LongAlpaca 实例结合，以缓解微调过程中的灾难性遗忘问题。
• 在 Q、K、V、O 投影头以及嵌入层上应用 QLoRA，LoRA 秩设为 32，缩放因子 alpha 为 16，批量大小为 8，学习率采用线性衰减，初始值为 5e-5。
• 将 RoPE 旋转位置编码的基长从 500K 扩展至 200M，以支持训练期间的更长上下文。
• 使用单台 8xA800（80GB）GPU 机器进行一周期训练，借助梯度检查点和无模型并行，8 小时内完成。

Figure 1: The accuracy score of Llama-3-8B-Instruct-80K-QLoRA on Needle-In-A-HayStack task. The blue vertical line indicates the training length, i.e. 80K.

实验结果

研究问题

• RQ1 少量由强大 LLM（如 GPT-4）生成的合成数据，是否足以实现对小型 LLM 的高效且有效的长上下文微调？
• RQ2 微调后的 LLM 在其最大训练长度之外（如 80K 到 128K）的泛化能力有多强？
• RQ3 通过此方法扩展上下文长度是否会影响模型在短上下文任务上的性能？
• RQ4 该扩展模型在标准长上下文基准测试中，与原始模型及其他长上下文 LLM 相比表现如何？
• RQ5 该方法是否可通过增加计算资源实现更长上下文长度的扩展？

主要发现

• Llama-3-8B-Instruct-80K-QLoRA 模型在 80K 全长上下文的 Needle-In-A-Haystack 任务中达到 100% 准确率，并在未见位置上泛化至 128K，表现优异。
• 在 Topic Retrieval 基准测试中，该模型在所有评估上下文长度（5K 至 70K）下均保持 100% 准确率，显著优于原始 Llama-3-8B-Instruct（超过 9K 后即失败）。
• 在 LongBench 基准上，该模型平均得分达 47.19，显著优于原始 Llama-3-8B-Instruct（43.20）和 Llama-3-8B-Instruct-262K（43.73），在 Single-Doc、Multi-Doc 和 Summ 等多个任务中表现突出。
• 在 InfiniteBench 上，该模型在 LongBookQA 上得分为 30.92，在 LongBookSum 上得分为 14.73，优于原始 Llama-3-8B-Instruct（7.00 和 16.40），并在摘要任务上接近 GPT-4 性能。
• 在 MMLU 上，该模型在零样本评估中得分为 64.44，略低于原始 Llama-3-8B-Instruct（65.91），表明在短上下文能力上存在轻微权衡，但仍优于其他同规模开源模型。
• 整个训练过程仅在单台 8xA800（80GB）GPU 上耗时 8 小时，展现出极高的效率与可复现性。

Figure 2: The accuracy of Topic Retrieval task.