Megalodon：针对大规模预训练的Mega改进

Megalodon是针对作者两年前提出的Mega架构改进版本，旨在提高大规模长上下文的预训练和推理效率。针对于Mega架构的几个局限性，Megalodon进一步提出了多个新技术来提高其学习能力和训练稳定性。在与Llama2模型完全相同的2T的预训练数据上，作者训练了一个7B大小的模型，使其可以和Llama2-7B进行完全公平的比较。最终取得了比Llama2更好的trainingloss。同时，在多个benchmark上的表现也明显好于同样大小的Llama2。

论文题目： Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length 论文链接： https://arxiv.org/abs/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上取得了很好的结果，但它依然存在一些问题：

1. 
   

   Mega w. chunk-attention的结果落后于Mega w. full-attention。Mega使用了chunk attetnion来将计算复杂度从平方降低到线性级别。但是从Mega论文中的结果可以看到，chunk-attention版本的Mega的效果还是明显落后于full-attention的版本。我们认为这主要是因为Mega中的EMA这个模块不够强大，无法很好的弥补chunk attention带来的上下文信息损失。

   
   

2. 
   

   模型架构对于不同数据和任务的差异。首先，这不是Mega独有的问题，而是几乎所有模型框架都存在的一个问题，即对于不同的数据类型和任务，虽然模型架构整体上差不多（比如都是Transformer），但实际还是存在很多架构的差异的。比如在LRA的不同任务上，大家会使用不同的normalization方法（BatchNorm, LayerNorm, RMSNorm, ScaleNorm）以及不同的normalization配置（Pre-Norm, Post-Norm, QK-Norm）。

   
   

3. 
   

   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将复杂度降低到了线性级别。

3.1 长上下文测试结果

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

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

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

四、结语

这个工作从有想法到最终完成，经历了近两年的时间。期间经历数次失败，也学习到了很多大规模预训练时代正确做科研的方法。得益于和Meta AI的合作，我们有相对充足的计算资源，以及可以和Llama2在完全相同数据上比较的机会。但也难免要受限于Meta内部的一些审核。比如因为审核我们错过了COLM的投稿，而且部分代码没办法放出来。训练好的7B模型也还在审核当中，什么时候能放出也不得而知。但是我们尽量放出了我们的代码，包括全部模块高效的cuda实现，以及并行训练所需要的全部技术。我们在github上提供一段预训练的伪代码。大家可以通过我们提供的代码来搭建一个预训练的脚本。

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

1. 
   

   对于两个不同模型架构的比较必须要在数据完全相同的条件下才有说服力。当数据不同的时候，哪怕不同的比例很小（<10%)，最后的结果也可能有明显的差别。包括training loss和下游任务的结果，都受到训练数据的很大影响。

   
   

2. 
   

   对于不同的架构，一定要在模型得到充分训练的条件下的比较才有意义。例如对于7B大小的模型，2T的训练数据几乎是基本要求。有的模型可能在数据少的时候表现的很好，但是数据规模增大后反而落后其他模型。因此，对于大模型架构的比较，结果有说服力的前提是充分的训练。

   
   

3. 
   

   对于架构差别很大的模型，传统的基于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.

作者：马学喆 来源：【知乎】https://zhuanlan.zhihu.com/p/692682649

llustration From IconScout By Delesign Graphics

-The End-

扫码观看！

本周上新！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（

www.techbeat.net

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

chenhongyuan@thejiangmen.com

或添加工作人员微信（chemn493）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。

关于我“门”

将门是一家以专注于数智核心科技领域的新型创投机构，也是北京市标杆型孵化器。 公司致力于通过连接技术与商业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈