Meta等发布巨齿鲨Megalodon！无限上下文长度的高效大模型预训练和推理

公众号新闻

2024-05-20 13:05

©作者 | Xuezhe Ma

单位 | 南加州大学

研究方向 | NLP, ML

论文标题：

Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length

论文地址：

https://arxiv.org/pdf/2404.08801

论文源码：

https://github.com/XuezheMax/megalodon

摘要

本文介绍了 Megalodon 架构。Megalodon 是针对我们两年前提出的 Mega 架构 [1] 的改进版本，旨在提高大规模长上下文的预训练和推理效率。针对于 Mega 架构的几个局限性，Megalodon 进一步提出了多个新的技术来提高它的学习能力和训练稳定性，包括 Complex Exponential Moving Average（CEMA）, Timestep Normalization, Normalized Attention mechanism 以及 Pre-Norm with two-hop residual configuration.

在与 Llama2 模型完全相同的 2T 的预训练数据上，我们训练了一个 7B 大小的 Megalodon 模型，使得我们可以和 Llama2-7B 进行完全公平的比较。最终，Megalodon 取得了比 Llama2 更好的 training loss。同时，在多个 benchmark 上 Megalodon 的表现也明显好于同样大小的 Llama2（详见下面的图和表格）

Megalodon架构

Mega 架构虽然在一些小规模的 benchmark 上取得了很好的结果，但它依然存在一些问题：

Mega w. chunk-attention 的结果落后于 Mega w. full-attention. Mega 使用了 chunk attetnion 来将计算复杂度从平方降低到线性级别。但是从 Mega 论文中的结果可以看到，chunk-attention 版本的 Mega 的效果还是明显落后于 full-attention 的版本。我们认为这主要是因为 Mega 中的 EMA 这个模块不够强大，无法很好的弥补 chunk attention 带来的上下文信息损失。
模型架构对于不同数据和任务的差异. 首先，这不是 Mega 独有的问题，而是几乎所有模型框架都存在的一个问题，即对于不同的数据类型和任务，虽然模型架构整体上差不多（比如都是 Transformer），但实际还是存在很多架构的差异的。比如在 LRA 的不同任务上，大家会使用不同的 normalization 方法（BatchNorm, LayerNorm, RMSNorm, ScaleNorm）以及不同的 normalization 配置（Pre-Norm, Post-Norm, QK-Norm）。
Mega 在大规模预训练上缺少证据. Mega 之前都是在小规模的数据和任务上进行的测试，并没有应用到大规模预训练上。

本文提出的 Megalodon 架构就是针对 Mega 以上的问题做改进。

2.1 Complex Exponential Moving Average

为了提高 Mega 中 EMA 模块的能力，直接的想法是将 EMA 扩展到复数域。这一想法是受到 SSM 中方法的启发。简单来说，我们通过添加一个极坐标的参数，可以直接将实数域的 EMA 扩展到复数域：

其中，为了让每个维度的 argument 能一致的分散在整个频域中，我们用如下的方法：

2.2 Timestep Normalization

在 Megalodon 中我们将 Group Normalization 扩展到 auto-regressive 的模式下，提出了 Timestep Normalization。

如图所示，在 Timestep Normalization 中，每一步的 mean 和 variance 是根据前面所有位置而累积计算的。为了能在 GPU 上面快速计算，我们提供了高效的 cuda 实现。

2.3 Normalized Attention in Megalodon

为了提高架构在大规模训练下的稳定性，我们针对 Mega 里面的 gated attention 模块提出了 normalized attention：

我们发现对Z做了 normalization 操作可以明显提高稳定性，同时传统 attention 中的 scale 项（一般为 sqrt(d)）也可以省去。

2.4 Pre-Norm with Two-hop Residual

为了进一步提高稳定性，我们提出了一种改进 Pre-Norm 中 residual connection 配置的新方案。通过简单的修改 residual connection，我们可以明显提高训练稳定性：

实验结果

为了公平的与 Llama2 进行比较，我们用了和它完全一样的 2T 的训练数据（得益于和 Meta AI 的合作）。我们分别比较了数据学习效率（data efficiency）和计算效率（computation efficiency）。

数据学习效率（data efficiency）简单来说就是不同的模型在相同的大小（model size）和相同的数据上进行训练，最终比较学习的效果。通过文章开始的 training loss 的图以及多个 benchmark 上面的结果可以看出，Megalodon 比 Transformer 在 7B 这个规模下有更好的数据学习效率。

计算效率（computation efficiency）是在相同的计算资源下比较相同大小的模型在训练中的速度。这里我们比较了 Megalodon 和 Transformer 在不同的上下文长度（4K vs. 32K）下的训练时每秒处理 token 的个数（token per second）。对于 Transformer，我们使用了 FlashAttention V2 来加速计算。

可以看到，当上下文长度比较短时 Transformer 比 Megalodon 快大约 5%。这是由于 Megalodon 要额外计算 CEMA 等模块。但是当上下文增加到 32k 的时候，Megalodon 比 Transformer 快 32%，这是得益于 chunk attention 将复杂度降低到了线性级别。

长上下文测试结果

为了测试 Megalodon 在长上下文上的泛化能力，我们首先测试了它在不同长度下的 PPL

可以看到，在我们将上下文长度从 4K 增加到 2M 的过程中，Megalodon 的 PPL 一直稳定下降。

我们还测试了 Megalodon 在 Scrolls 上的结果，同时也测试了一些中小规模的任务，比如 LRA，WikiText-103, PG19, ImageNet-1K 以及 Speech Commands。具体结果请阅读论文。

结语

这个工作从有想法到最终完成，经历了近两年的时间。期间经历数次失败，也学习到了很多大规模预训练时代正确做科研的方法。得益于和 Meta AI 的合作，我们有相对充足的计算资源，以及可以和 Llama2 在完全相同数据上比较的机会。但也难免要受限于 Meta 内部的一些审核。比如因为审核我们错过了 COLM 的投稿，而且部分代码没办法放出来。

训练好的 7B 模型也还在审核当中，什么时候能放出也不得而知。但是我们尽量放出了我们的代码，包括全部模块高效的cuda实现，以及并行训练所需要的全部技术。我们在 github 上提供一段预训练的伪代码。大家可以通过我们提供的代码来搭建一个预训练的脚本。

通过这个项目，我们也体会到了在大模型时代做新的模型架构时要注意的问题。总结来说：

对于两个不同模型架构的比较必须要在数据完全相同的条件下才有说服力。当数据不同的时候，哪怕不同的比例很小（<10%），最后的结果也可能有明显的差别。包括 training loss 和下游任务的结果，都受到训练数据的很大影响。
对于不同的架构，一定要在模型得到充分训练的条件下的比较才有意义。例如对于 7B 大小的模型，2T 的训练数据几乎是基本要求。有的模型可能在数据少的时候表现的很好，但是数据规模增大后反而落后其他模型。因此，对于大模型架构的比较，结果有说服力的前提是充分的训练。
对于架构差别很大的模型，传统的基于 flops 的 scaling law 的比较意义在降低。原因是两个不同架构的模型，即使有相同的 flops，他们的实际速度可能差几倍。这个和架构算法本身是不是适合在最先进的 GPU 上计算有很大的关系。因此，真正贴合实际的比较方法是像本文中那样分成数据学习效率和计算效率两个方面。但是这样在实际中对于研究员的工程能力有很高的要求。在大模型时代，新算法的开发已经和系统等方面高度结合在一起。

参考文献

[1] Xuezhe Ma and Chunting Zhou and Xiang Kong and Junxian He and Liangke Gui and Graham Neubig and Jonathan May and Luke Zettlemoyer. Mega: Moving Average Equipped Gated Attention. ICLR-2023.

更多阅读