分享

终结长文本之战:Meta爆改Transformer,实现无限上下文长度的新算法

 长沙7喜 2024-04-17 发布于湖南

图片

李然、陈斯达

编辑苏建勋

封面来源企业官网

大模型上下文长度作为影响大模型应用潜力的一个非常重要的指标,各家大模型都在用力卷。

国外有谷歌Gemini Pro 1.5号称支持200万到1000万token的长上下文,国内的Kimi也是将长上下文作为通往AGI的最关键要素。

而广大的大模型研究人员,同样在为了让Transformer能够高效地增加上下文长度而头秃。

但是由于Transformer本身机制的问题,上下文长度的增加一定会让训练和推理的成本大幅增加。

于是,来自Meta,南加州,CMU,UCSD的研究人员决定把桌子掀了,直接从第一性原理出发,以无限长上下文为目标重做一个新的“Transformer”。

他们提出的Megalodon,是一种能够高效处理具有无限上下文长度序列的新型神经网络架构。

图片

论文地址:https:///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都表现显著优势。

图片

作者介绍

Xuezhe (Max) Ma
图片
他本科硕士毕业于上海交通大学,博士毕业于CMU,现在在CMU任职研究助理。
Beidi Chen
图片

她本科毕业于UC伯克利,博士毕业于莱斯大学,现在为Meta FAIR的客座科学家,同时也是CMU的助理教授。

Wenhan Xiong
图片
本科毕业于中国科学技术大学,博士UCSB,目前在Meta做生成式AI方面的研究。
Lili Yu
本科毕业于北大,博士毕业于MIT,曾就职于AI初创企业Asapp,现为Facebook AI研究员。

    本站是提供个人知识管理的网络存储空间,所有内容均由用户发布,不代表本站观点。请注意甄别内容中的联系方式、诱导购买等信息,谨防诈骗。如发现有害或侵权内容,请点击一键举报。
    转藏 分享 献花(0

    0条评论

    发表

    请遵守用户 评论公约

    类似文章 更多