巨齿龙：无限上下文长度的 LLM 预训练与推理

近年来，大型语言模型（LLM）在各个领域取得了显著的进步，但其在处理长序列数据方面仍然存在挑战。传统的 Transformer 架构由于其二次复杂度和对长度泛化的有限归纳偏差，在处理长序列时效率低下。为了克服这些限制，研究人员提出了各种亚二次复杂度解决方案，例如线性注意力和状态空间模型，但这些方法在预训练效率和下游任务准确性方面通常不如 Transformer。

本文介绍了 巨齿龙（MEGALODON），一种能够高效处理无限上下文长度的序列建模神经架构。巨齿龙继承了 MEGA（带门控注意力的指数移动平均）的架构，并引入了多个技术组件来提升其能力和稳定性，包括：

复指数移动平均（CEMA）：将 MEGA 中的多维衰减 EMA 扩展到复数域，进一步提升了 EMA 的表达能力。
时间步归一化层：将组归一化层推广到自回归序列建模任务，允许沿序列维度进行归一化，从而更好地处理序列数据中的内部协变量偏移。
归一化注意力机制：通过对共享表示进行归一化，提高了注意力机制的稳定性，并简化了模型结构。
带两跳残差的预归一化：通过重新排列每个块中的残差连接，有效地缓解了预归一化在模型规模扩大时出现的稳定性问题。

巨齿龙通过将输入序列分割成固定长度的块，在模型训练和推理过程中都实现了线性计算和内存复杂度。

巨齿龙的优势

在与 LLAMA2 的对比实验中，巨齿龙在 70 亿参数规模和 2 万亿训练数据的情况下，展现出比 Transformer 更高的效率。巨齿龙的训练损失达到了 1.70，介于 LLAMA2-7B（1.75）和 LLAMA2-13B（1.67）之间。

巨齿龙在各种下游任务和数据模态的基准测试中，都表现出比 Transformer 更强的鲁棒性，包括：

常识推理：HellaSwag、PIQA、SIQA、WinoGrande、ARC-e 和 ARC-c
世界知识：NaturalQuestions 和 TriviaQA
阅读理解：BoolQ
综合评估：MMLU
长上下文问答：Scrolls 数据集中的 NarrativeQA、Qasper 和 QMSum
指令微调：MT-Bench
图像分类：ImageNet-1K
自回归语言建模：PG-19

巨齿龙的未来展望

巨齿龙的成功为未来的 LLM 研究指明了方向：

巨齿龙能够高效地处理无限上下文长度的序列数据，为处理长文本、多轮对话和视频生成等任务提供了新的可能性。
巨齿龙在各种数据模态的基准测试中都表现出色，表明它有潜力成为未来大规模多模态预训练的通用架构。

参考文献

Ma, X., Yang, X., Xiong, W., Chen, B., Yu, L., Zhang, H., … & Zhou, C. (2023). MEGALODON: Efficient LLM Pretraining and Inference with Unlimited Context Length. arXiv preprint arXiv:2404.08801.

https://gateway.ipfsscan.io/ipfs/QmY8ia4azH6H351Ss1rXjg27idXyDwdpKWR1kBrrrUd43Q?filename=从第一性原理出发，让深度学习更加生动.pdf