继Mamba之后,又一敢于挑战Transformer的架构诞生了!
来自Meta、南加州大学(USC)、CMU和UCSD的研究人员提出了全新的神经网络架构——Megalodon(巨齿鲨)。
这是专为有效处理「无限上下文」长度的LLM预训练,以及推理而设计的架构。
论文地址:
我们都知道,Transformer架构个在处理长上下文时,会受到二次复杂度,以及长度外推能力弱的限制。
尽管已有次二次方解决方案(诸如线性注意力,状态空间模型),但它们在预训练效率,甚至下游任务的准确率上,通常还不及Transformer。
Megalodon的出现,就是为了解决无限处理上下文的难题。
同时,它可以同时实现高效训练(减少通信和计算量),以及高效推理(保持恒定的KV缓存)。
值得一提的是,在与Llama 2的直接比较中,Megalodon在处理70亿参数和2万亿训练token的任务上,不仅训练更高效,而且准确率也超过了Transformer。
具体来说,Megalodon的训练损失为1.70,位于Llama2-7B(1.75)和 13B(1.67)之间。
这一改变范式的创新代表着AI领域的巨大飞跃,Megalodon开启了计算效率和性能的新时代。
GPT-3发布以来最大里程碑
网友表示,先是谷歌,又是Meta,无限上下文离我们更进一步,LLM将会释放出无限潜力。
还有人认为「无限上下文长度,绝对是游戏规则的改变者」!
更有甚者,初创公司CEO称,「这是自GPT-3发布以来最大的里程碑,但却没有任何动静?!
Megalodon就相当于是AGI的基础」。
「Meta的Megalodon是一项突破性进展,对AGI具有重要意义。它的无限上下文长度模拟了人类的认知,实现了无缝任务切换」。
论文作者Hao Zhang表示,这是一种全新替代Transformer的架构。
论文作者Beidi Chen称,「注意力虽好,但你不需要完整的注意力机制」!
普林斯顿助理教授Tri Dao表示,「将SSM/RNN/EMA与注意力相结合是获得更高质量、更长上下文和更快推理的方法!Griffin、Jamba、Zamba和现在的Megalodon都是很好的例子」。
革命性架构,训练更稳定
那么,Megalodon架构采用了怎样的设计,才能取得如此优异的表现?
据介绍,它基于MEGA架构进行了改进,并新增了多个技术组件。
首先,复杂指数移动平均(CEMA)组件是一种全新技术,扩展了MEGA中使用的多维阻尼指数移动平均方法到复数域,可以增强模型处理复杂数据的能力。
其次,研究人员提出了一种创新的归一化技术——「时间步归一化层」。
它将传统的组归一化技术扩展到自回归序列建模任务中,允许模型在处理序列数据时,进行有效的归一化。
以往,「层归一化」(Layer Normalization)与Transformer相结合性能,虽令人印象深刻。
但很明显,层归一化并不能直接减少时间步长或顺序维度的内部协变量偏移。
另外,「组归一化」(Group Normalization)虽比「层归一化」在CV任务中获得改进,但它却无法直接应用于Transformer的自回归序列建模,因未来信息会通过时间步维度的均值和方差泄漏。
如下图所示,c展示了Megalodon架构中,层标准化和时间步标准化的方法。
最后,研究人员为了增强大规模LLM预训练的稳定性,提出了将归一化注意力,和带有两跳残差的预归一化相结合的配置。
这种配置可以优化模型的学习过程,提高训练的稳定性。
下图3中,a是Megalodon的完整框架草图。
中间和右边两张图分别介绍了,预归一化和带有两跳残差预归一化的配置。
2T token训练,性能超越Llama2-7B
在具体实验评估中,研究人员将Megalodon扩展到70亿参数规模,并将其应用于2万亿token的大规模LLM预训练中。
此外,作者还在中/小参数规模的序列建模基准上进行了实验,包括Long Range Arena (LRA) 、Speech Commands上的原始语音分类、ImageNet-1K上的图像分类,以及WikiText-103和PG19上的语言建模。
结果显示,在这些任务中,Megalodon在各种数据模式下的表现明显优于所有最先进的基线模型。
数据学习效率
通过训练损失图以及多个benchmark的结果可以看出,Megalodon比Transformer在7B参数下有更好的数据学习效率。
计算效率
针对不同的4K和32K上下文长度,Megalodon这一架构的预训练的计算效率也是非常强的。
学术基准上短上下文评估
具体来说,研究人员在短上下文(4K token)的标准学术基准上,对Megalodon与Llama 2,以及开源基础模型进行了比较。
在相同的2万亿token训练后,Megalodon-7B的表现明显优于Llama2-7B。
长上下文评估
针对不同长上下文困惑度,证明了Megalodon可以利用很长的上下文进行下一个token预测的能力。
图5显示了,验证数据集在4K到2M各种上下文长度下的困惑度(PPL)。
在Scroll数据集中的长上下文QA任务中,Megalodon在NaQA上获得最佳F1,并与Llama 2 Long相竞争。
中等规模基准评估
在Long Range Arena(LRA)的测试中,新架构显著缩小了分块注意力和全注意力之间的性能差距。
其他评测集,如原始语音分类、ImageNet-1K、WikiText-103和PG-19的结果如下:
一些感想
这里quote一下这项研究原作者的一些感悟和经历:
通过这个项目,研究者们也体会到了在大模型时代做新的模型架构时要注意的问题。总结来说:
原文链接: