终结长文本之战：Meta爆改Transformer，实现无限上下文长度的新算法
文｜李然、陈斯达
编辑｜苏建勋
大模型上下文长度作为影响大模型应用潜力的一个非常重要的指标，各家大模型都在用力卷。
国外有谷歌Gemini Pro 1.5号称支持200万到1000万token的长上下文，国内的Kimi也是将长上下文作为通往AGI的最关键要素。
而广大的大模型研究人员，同样在为了让Transformer能够高效地增加上下文长度而头秃。
但是由于Transformer本身机制的问题，上下文长度的增加一定会让训练和推理的成本大幅增加。
于是，来自Meta，南加州，CMU，UCSD的研究人员决定把桌子掀了，直接从第一性原理出发，以无限长上下文为目标重做一个新的“Transformer”。
他们提出的Megalodon，是一种能够高效处理具有无限上下文长度序列的新型神经网络架构。
论文地址：https://arxiv.org/abs/2404.08801
Megalodon不仅继承了MEGA构架的优点，并且引入了多个创新组件，例如复杂指数移动平均（Complex Exponential Moving Average, CEMA）、时间步归一化层（timestep normalization layer）、归一化注意力机制（normalization attention）以及预归一化和双跳残差配置，这些改进提高了模型的能力和稳定性。
通过Megalodon，研究人员实现了高效训练（减少通信和计算）和低成本高效推理（恒定KV缓存）的两全其美。
在与Llama2的对照正面比较中，Megalodon在7B和2T训练标记的规模上比Transformer实现了更好的效率和准确性。
训练效率比Llama2高出了不少。
在各种学术基准测试中，Megalodon的表现也优于Llama2。
Megalodon的PPL（7b，训练上下文32K）在验证集上从4K到2M的各种上下文长度上的表现。
Scrolls中长上下文QA任务的结果：
研究人员通过引入了复指数移动平均（complex exponential moving average，CEMA）成分，它将MEGA中的多维阻尼（multi-dimensional damped）EMA扩展到了复数域（complex domain）。
然后他们还提出了时间步归一化层（timestep normalization layer,），它将组归一化层推广到自动回归序列建模任务中，允许沿序列维度进行归一化。
为了提高大规模预训练的稳定性，他们还提出归一化注意力和双跳残差配置预归一化：
Megalodon在语音分类、ImageNet-1K、WikiText-103和PG19的性能如下：
在Long Range Arena（LRA）上的表现也非常好，大大缩小了块注意力与完全注意力之间的差距。
作者表示，只要Meta允许，他们将很快放出模型供大家体验。
网友看了论文之后表示，画面感太强了。
这个无限上下文的“Transformer”，会是人类通往AGI的捷径吗？
基于MEGA优化架构，更高效更稳定
Megalodon的架构到底好在哪？
简单来说，就是在MEGA的基础上，新加了几个技术组件——复指数移动平均(CEMA)、时间步归一、归一化注意力以及双跳残差连接预归一化。
CEMA，将MEGA中的多维阻尼EMA扩展到复数域，让更丰富的上下文建模成为可能，让模型处理复杂数据时更加强大。
之前的层归一化内部协变量移位比较高。时间步归一化将前一代的组归一化推广到自回归建模中，降低时间维度上的内部协变量移位。这一关键功能可让Megalodon实现大规模数据集的无缝处理。
稳定训练深度架构中，归一化配置至关重要，而预归一化已经成为默认配置。然而，扩大模型规模时，预归一化可能会不稳定。
△图中为Megalodon的完整框架、预归一化、带有两跳残差预归一化的配置
此次Megalodon中加入的双跳残差连接预归一化，正是缓解了扩大模型规模时，预归一化的不稳定“症状”。
Megalodon的成功算是说明了，模型创新，离不开对既有框架进行增量迭代优化。
性能优于Llama 2，文本越长理解越到位
为评估在长文本序列建模方面的效率，研究将Megalodon扩展到70亿参数规模，并应用于2万亿token的大语言模型训练。各种任务一路比下来，Megalodon明显优于很多最先进的基准模型。
数据与运算效率
基于相同计算资源，以每秒生成的token数量为基准，研究比较LLAMA2-7B和Megalodon-7B在不同文本长度条件下的训练速度。
在4K上下文长度下，Megalodon-7B稍慢于LLAMA2-7B。但当上下文长度扩展到32K时，Megalodon-7B比LLAMA2-7B快得多。这无疑展示了Megalodon在长文本训练中的运算效率。
短上下文评测
短文本评测中采用标准化测试集合，包括常识推理、世界知识、阅读理解等。在所有基准测试中，Megalodon-7B都超过了LLAMA2-7B。但研究也说，不能与Mistral-7B和Gemma-8B等其他开源模型直接比——人家训练数据集比Megalodon-7B用的要大很多。
长上下文评测
随着文本长度的增加，Megalodon处理长序列数据时计算的困惑度（Perplexity over Long Sequences），呈单调减少。这也直观显示出，Megalodon在理解超长序列数据时有效且稳健。
研究还对Megalodon进行了长文本“开卷考试”。参考的是Scrolls数据集中的任务，包括NarrativeQA、Qasper和QMSum。
在与其他7B规模的开源模型比较中，Megalodon-7B在NarrativeQA上获得最佳分数。与LLAMA2-7B Long相比，在另外两项任务中，成绩也不分上下。
指令微调
为了评估Megalodon在跟随指令和对齐方面的泛化能力，研究在受控环境下基于专门的指令对齐数据，对Megalodon-7B的基础模型展开微调。过程中，没有用人类反馈进行强化学习（RLHF）。
在名为MT-Bench的多轮对话基准测试中，同样是7B规模，与Vicuna相比，Megalodon性能更加优越。与利用RLHF的LLAMA2-Chat相比，也水平相当。
中等规模基准评测
研究在Imagenet-1K数据集上进行实验，发现在图像分类任务的表现上，相比于DeiT-B，Megalodon准确率率比DeiTy-B提高约1.3%，比MEGA提高0.8%。
研究评估了Megalodon在中等规模PG19数据集上的自回归语言建模，相比基线模型，Megalodon都表现显著优势。
作者介绍
她本科毕业于UC伯克利，博士毕业于莱斯大学，现在为Meta FAIR的客座科学家，同时也是CMU的助理教授。
36氪旗下AI公众号