作者: Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, Jiaya Jia
提交/修订日期: 2023年9月21日提交，2024年3月8日修订 (v3)
主题/分类: 计算与语言 (cs.CL); 人工智能 (cs.AI); 机器学习 (cs.LG)
本文提出了 LongLoRA，一种高效微调方法，旨在以有限的计算成本扩展预训练大语言模型 (LLMs) 的上下文长度。通常，使用长上下文训练 LLMs 计算成本高昂，需要大量的训练时间和 GPU 资源。例如，在上下文长度为 8192 上训练的自注意力层计算成本是 2048 的 16 倍。
本文从两个方面加速 LLMs 的上下文扩展：
1. 稀疏局部注意力微调：尽管推理时需要密集的全局注意力，但模型的微调可以有效地通过稀疏局部注意力完成。提出的移位稀疏注意力 (shifted sparse attention, S^2-Attn) 能有效实现上下文扩展，在性能与使用原始注意力微调相似的同时，带来显著的计算节省。该方法在训练中仅需两行代码即可实现，且在推理时是可选的。
2. 参数高效微调机制：我们重新审视了用于上下文扩展的参数高效微调机制。值得注意的是，我们发现 LoRA 在可训练的嵌入层和归一化层的前提下，对于上下文扩展效果良好。LongLoRA 将这种改进的 LoRA 与 S^2-Attn 相结合。
LongLoRA 在 Llama2 模型 (7B/13B 到 70B) 的各种任务上展示了强大的实证结果。LongLoRA 可以在单个 8x A100 机器上将 Llama2 7B 的上下文从 4k 扩展到 100k，或将 Llama2 70B 扩展到 32k。LongLoRA 在扩展模型上下文的同时保留了其原始架构，并且与大多数现有技术兼容，例如 Flash-Attention2。此外，我们还使用 LongLoRA 和我们构建的长指令跟随数据集 LongAlpaca 进行了监督微调。
代码与数据: 代码、模型、数据集和演示可在 GitHub 获取。