Meta无限长文本大模型来了：参数仅7B，已开源

谷歌之后，Meta 也加入了无限长上下文的竞赛。
Transformer 的二次复杂度和糟糕的长度外推能力，让它们在处理长序列时力不从心。虽然也有线性注意力、状态空间模型这些次二次方案，但从实际效果看，预训练效率和下游任务准确率都差点意思。

长文本处理一直是大型语言模型的老大难问题。前几天谷歌拿出了 Infini-Transformer，号称能把基于 Transformer 的 LLM 扩展到无限长输入，还不增加内存和计算开销，一下子就火了。几乎同一时间，Meta 也亮出了自己的无限长文本技术。

• 论文地址：https://arxiv.org/pdf/2404.08801.pdf
• 论文标题：MEGALODON: Efficient LLM Pretraining and Inference with Unlimited Context Length
• 代码：https://github.com/XuezheMax/megalodon

在 4 月 12 日提交的这篇论文中，来自 Meta、南加州大学、CMU、UCSD 等机构的研究者推出了 MEGALODON——一个上下文长度不受限制的高效序列建模架构。它继承了 MEGA（带门控注意力的指数移动平均）的框架，并引入了几项关键技术来提升能力和稳定性，包括复数指数移动平均（CEMA）、时间步归一化层、归一化注意力机制，以及带两跳残差的预归一化配置。

拿 MEGALODON 和 LLAMA2 直接比：70 亿参数、2 万亿训练 token 的规模下，它的效率比 Transformer 更好。训练损失达到了 1.70，正好落在 LLAMA2-7B（1.75）和 13B（1.67）之间。在一系列跨任务、跨模态的基准测试中，MEGALODON 相对 Transformer 的改进都相当显著。

MEGALODON 本质上是 MEGA（Ma et al., 2023）的改良版，利用了门控注意力机制和经典的指数移动平均（EMA）方法。为了进一步提升大规模长上下文预训练的能力和效率，作者们带来了几个新组件。第一，引入复数指数移动平均（CEMA），把 MEGA 中的多维阻尼 EMA 扩展到了复数域。第二，提出时间步长归一化层，把组归一化推广到自回归序列建模场景，允许沿着序列维度做归一化。

为了提高大规模预训练的稳定性，MEGALODON 还设计了归一化注意力，以及一种带两跳残差配置的预归一化——这是对广泛使用的预归一化和后归一化方法的改造。只要简单地把输入序列切成固定大小的块（就像 MEGA-chunk 那样），训练和推理时就能实现线性的计算和内存复杂度。

在和 LLAMA2 的直接对比中，控制数据和计算量相同的情况下，MEGALODON-7B 的训练困惑度显著优于训练 LLAMA2-7B 所用的最先进 Transformer 变体。长上下文建模评估——包括从 4K 到 2M 各种上下文长度的困惑度，以及 Scrolls 中的长上下文 QA 任务——都能证明 MEGALODON 建模无限长度序列的能力。再加上 LRA、ImageNet、Speech Commands、WikiText-103 和 PG19 等中小型基准的实验结果，MEGALODON 在规模和多模态上的实力得到了全面验证。

方法介绍

先来快速回顾一下 MEGA（带门控注意力的移动平均）架构的关键组件，以及它存在的一些问题。

MEGA 把 EMA（指数移动平均）组件嵌入到注意力矩阵的计算中，从而引入了跨时间步维度的归纳偏置。具体来说，多维阻尼 EMA 首先通过扩展矩阵将输入序列 X 的每个维度单独扩展到 h 维，然后在 h 维隐藏空间中应用阻尼 EMA。形式如下：

为了降低全注意力机制的二次复杂度，MEGA 简单地把查询、键和值序列拆分为长度为 c 的块，注意力单独在每个块内计算，从而得到线性复杂度 O(kc²)=O(nc)。

从技术上讲，MEGA 中的 EMA 子层有助于捕获每个 token 附近的局部上下文，缓解了跨块边界丢失信息的问题。尽管 MEGA 表现令人印象深刻，但存在几个问题：

i) EMA 子层的表达能力有限，导致带块级注意力的 MEGA 性能仍落后于全注意力版本；
ii) 对于不同任务和数据类型，最终的 MEGA 架构可能存在差异，比如归一化层、归一化模式、注意力函数 f(·) 都不一样；
iii) 没有经验证据表明 MEGA 能扩展到大规模预训练。

CEMA：将多维阻尼 EMA 扩展到复数域

为了解决 MEGA 的问题，研究者开发了 MEGALODON。他们创造性地提出了复数指数移动平均（CEMA），把原来的公式改写成了复数形式：

（此处为原公式改写示意，实际保留原文图片）

并且将参数 θ_j 参数化为：

（原图）

时间步（Timestep）归一化

层归一化与 Transformer 结合的表现虽然惊艳，但很明显，它不能直接减少沿空间维度（时间步或序列维度）的内部协变量偏移。MEGALODON 通过计算累积均值和方差，把组归一化扩展到了自回归场景：

（原图）

图 2 展示了层归一化和时间步标准化的区别。

MEGALODON 中的归一化注意力

此外，研究者专门为 MEGA 定制了一种归一化注意力机制，以提高稳定性。具体形式如下：

（原图）

那么上文中的注意力操作就变成了：

（原图）

具有 Two-hop 残差的预范数（Pre-Norm）

调查发现，扩大模型规模会导致预归一化不稳定问题。基于 Transformer 块的预归一化可以表示为（如图 3(b) 所示）：

（原图）

在原始 MEGA 架构中，用 φ 来缓解这一问题（门控残差连接）。但更新门 φ 引入了更多参数，模型规模扩大到 70 亿时，不稳定问题仍然存在。MEGALODON 引入了一种名为 pre-norm 的新配置，带有 two-hop 残差，只是重新排列了每个块中的残差连接，如图 3(c) 所示：

（原图）

实验

为了评估 MEGALODON 在长上下文序列建模上的可扩展性和效率，研究者把它扩展到了 70 亿参数规模。

LLM 预训练

为了提高数据效率，训练过程中展示了 MEGALODON-7B、LLAMA2-7B 和 LLAMA2-13B 的负对数似然（NLL），如图 1 所示。在相同数量的训练 token 下，MEGALODON-7B 获得了比 LLAMA2-7B 明显更低的 NLL，表现出更好的数据效率。

（图 1）

图 4 说明了分别使用 4K 和 32K 上下文长度时，LLAMA2-7B 和 MEGALODON-7B 在每个设备上的平均 WPS（每秒词/ token）。对于 LLAMA2 模型，使用了 Flash-Attention V2 加速全注意力计算。在 4K 上下文长度下，由于引入了 CEMA 和时间步归一化，MEGALODON-7B 比 LLAMA2-7B 稍慢约 6%。当扩展到 32K 上下文长度时，MEGALODON-7B 明显更快（约 32%），证明了它对长上下文预训练的计算效率。

（图 4）

短上下文评估

表 1 总结了 MEGALODON 和 LLAMA2 在学术基准上的结果，同时与 MPT、RWKV、Mamba、Mistral、Gemma 等其他开源基础模型进行了比较。在相同的 2T token 上预训练后，MEGALODON-7B 在所有基准测试中都优于 LLAMA2-7B。在某些任务上，它的性能甚至与 LLAMA2-13B 相当或更好。

（表 1）

长上下文评估

图 5 显示了验证数据集在 4K 到 2M 各种上下文长度下的困惑度（PPL）。可以看到 PPL 随着上下文长度单调下降，验证了 MEGALODON 在建模极长序列方面的有效性和鲁棒性。

（图 5）

指令微调

表 3 总结了 7B 模型在 MT-Bench 上的性能。与 Vicuna 相比，MEGALODON 表现更优，且与 LLAMA2-Chat 相当——后者可是用了 RLHF 进一步对齐微调的。

（表 3）

中等规模基准评估

为了评估 MEGALODON 在图像分类任务上的表现，研究者在 ImageNet-1K 上进行了实验。表 4 报告了验证集上的 Top-1 准确率。MEGALODON 的准确率比 DeiT-B 高 1.3%，比 MEGA 高 0.8%。

（表 4）

表 5 展示了 MEGALODON 在 PG-19 上的字级困惑度（PPL），以及与之前最先进模型（包括 Compressive Transformer、Perceiver AR、块循环 Transformer、MEGABYTE 等）的对比。MEGALODON 明显领先。

（表 5）