机器之心报道

编辑：泽南、陈萍

谷歌之后，Meta 也来卷无限长上下文。

Transformers 的二次复杂度和弱长度外推限制了它们扩展到长序列的能力，虽然存在线性注意力和状态空间模型等次二次解决方案，但从以往的经验来看，它们在预训练效率和下游任务准确性方面表现不佳。

长文本是大语言模型一直在努力的方向。近日，谷歌提出的 Infini-Transformer 引入有效方法，可以将基于 Transformer 的大型语言模型 (LLM) 扩展到无限长输入，而不增加内存和计算需求，吸引了人们的关注。

几乎就在同时，Meta 也提出了一种无限长文本技术。

在 4 月 12 日提交的一篇论文中，来自 Meta、南加州大学、CMU、UCSD 等公司、机构引入了 MEGALODON，一种用于高效序列建模的神经架构，上下文长度不受限制。

MEGALODON 继承了 MEGA（带有门控注意力的指数移动平均）的架构，并进一步引入了多种技术组件来提高其能力和稳定性，包括复数指数移动平均（CEMA）、时间步归一化层、归一化注意力机制和具有两个特征的预归一化（pre-norm）残差配置。

在与 LLAMA2 的直接比较中，MEGALODON 在 70 亿参数和 2 万亿训练 token 的规模上取得了比 Transformer 更好的效率。MEGALODON 的训练损失达到 1.70，处于 LLAMA2-7B (1.75) 和 13B (1.67) 之间。MEGALODON 相对于 Transformers 的改进在不同任务和模式的一系列基准测试中表现强劲。

MEGALODON 本质上是一种改进的 MEGA 架构（Ma et al., 2023），它利用了门控注意力机制和经典指数移动平均（EMA）方法。为了进一步提高 MEGALODON 在大规模长上下文预训练方面的能力和效率，作者提出了多种技术组件。首先，MEGALODON 引入了复数指数移动平均 (CEMA) 组件，将 MEGA 中的多维阻尼 EMA 扩展到复数域。其次，MEGALODON 提出了时间步长归一化层，它将组归一化层推广到自回归序列建模任务，以允许沿顺序维度进行归一化。

为了提高大规模预训练的稳定性，MEGALODON 进一步提出了归一化注意力，以及通过修改广泛采用的预归一化和后归一化方法而具有两跳残差配置的预归一化。通过简单地将输入序列分块为固定块，如 MEGA-chunk 中所做的那样，MEGALODON 在模型训练和推理中实现了线性计算和内存复杂性。

在与 LLAMA2 的直接比较上，控制了数据和计算的同时，MEGALODON-7B 在训练困惑度方面显著优于用于训练 LLAMA2-7B 的最先进的 Transformer 变体。在对长上下文建模的评估上，包括高达 2M 的各种上下文长度中的困惑度以及 Scrolls 中的长上下文 QA 任务证明了 MEGALODON 对无限长度序列进行建模的能力。中小型基准的其他实验结果，包括 LRA、ImageNet、Speech Commands、WikiText-103 和 PG19 证明了 MEGALODON 在体量和多模态上的能力。

方法介绍

首先，文章简单回顾了 MEGA（ Moving Average Equipped Gated Attention ）架构中的关键组件，并讨论了 MEGA 中存在的问题。

MEGA 将 EMA（ exponential moving average ） 组件嵌入到注意力矩阵的计算中，以纳入跨时间步长维度的归纳偏差。具体而言，多维阻尼 EMA 首先通过扩展矩阵 将输入序列 X 的每个维度单独扩展为 h 维，然后将阻尼 EMA 应用于 h 维隐藏空间。形式如下：

为了降低全注意力机制中的二次复杂度，MEGA 简单地将 (14-16) 中的查询、键和值序列拆分为长度为 c 的块。(17) 中的注意力单独应用于每个块，产生线性复杂度 O (kc^2 ) = O (nc)。

从技术上讲，MEGA 中的 EMA 子层有助于捕获每个 token 附近的本地上下文信息，从而缓解了在超出块边界的上下文中丢失信息的问题。尽管 MEGA 取得了令人深刻的印象，但面临如下问题：

i）由于 MEGA 中 EMA 子层的表达能力有限，具有块级注意力的 MEGA 性能仍然落后于全注意力 MEGA。

ii) 对于不同的任务、数据类型，最终的 MEGA 架构可能存在架构差异，例如，不同的归一化层、归一化模式和注意力函数 f (・) 。

iii) 没有经验证据表明 MEGA 可扩展用于大规模预训练。

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

为了解决 MEGA 面临的问题，该研究提出了 MEGALODON。

具体而言，他们创造性地提出了复指数移动平均 CEMA（ complex exponential moving average ），将上式（1）改写为如下形式：

并将（2）中的 θ_j 参数化为：

时间步（Timestep）归一化

尽管层归一化与 Transformer 相结合的性能令人印象深刻，但很明显，层归一化不能直接减少沿空间维度（也称为时间步长或序列维度）的内部协变量偏移。

在 MEGALODON 中，该研究通过计算累积均值和方差将组归一化扩展到自回归情况。

图 2 说明了层归一化和时间步标准化。

MEGALODON 中的归一化注意力

此外，该研究还提出了专门为 MEGA 定制的归一化注意力机制，以提高其稳定性。形式如下：

则上式 (17) 中的注意力操作改为：

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

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

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

实验

为了评估 MEGALODON 在长上下文序列建模上的可扩展性和效率，本文将 MEGALODON 扩展到 70 亿规模大小。

LLM 预训练

为了提高数据效率，研究者在训练过程中显示了 MEGALODON-7B、LLAMA2-7B 和 LLAMA2-13B 的负对数似然 (NLL)，如图 1 所示。

在相同数量的训练 token 下，MEGALODON-7B 获得了比 LLAMA2-7B 明显更好（更低）的 NLL，表现出更好的数据效率。

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

短上下文评估

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

长上下文评估

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

指令微调

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

中等规模基准评估

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

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

更详细内容请参考论文原文。