MiniMax豪赌线性注意力:钟毅然访谈

MiniMax豪赌线性注意力:MiniMax-01架构负责人钟毅然访谈

尽管 Transformer 架构目前主导着生成式 AI 领域,但它并非没有局限性,挑战者正在涌现。MiniMax-01 大胆拥抱线性注意力机制,并扩展到前所未有的 4560 亿参数,是开源社区中的一个颠覆者。这既是一场技术赌博,也是架构创新中的一个潜在里程碑。

在这次访谈中,我们与 MiniMax-01 架构负责人钟毅然进行了对话,探讨了线性注意力从实验室走向工业级大型模型的历程,以及他对模型架构的思考和见解。

一位非主流技术路线的先驱

您能简单介绍一下自己吗?

我是钟毅然,MiniMax 的高级研究总监,主要负责网络架构设计和多模态理解大型模型。在 MiniMax,我的主要职责是领导 MiniMax-01 网络结构的设计。

之前,我曾在上海人工智能实验室担任新架构探索小组的负责人,专注于非 transformer 架构的高效训练建模方法以及视觉-音频-语言多模态融合的研究。

您是什么时候开始研究线性注意力的,为什么选择了这条技术路线?

我大约在 2021 年 7 月开始研究线性注意力。这源于我 2020 年博士期间发表的一篇论文,《可逆注意力》(Invertible Attention)。当时,可逆神经网络和注意力机制都非常流行,所以我们将它们结合起来进行了研究。

后来,我们团队的一些成员对数学非常感兴趣。像线性注意力这样的高效序列建模方法需要强大的数学基础,并涉及大量的公式推导,这与团队的兴趣非常吻合,所以我们选择了这个方向。

当时线性注意力在行业内的状况如何?

非常非主流,很少有人研究。大多数研究人员都专注于 transformers,它们实际上已经成为 NLP 领域的主导力量。

我们认为,与其成为人群中又一张做 transformer 研究的面孔,不如做一些不同的事情。

您如何评估线性注意力路线的技术潜力?

我们最初的动机很简单:解决 transformers 的二次计算复杂度问题。我们测试了各种方法,包括稀疏 transformers 和线性注意力。

我们发现稀疏 transformers 确实有效,与 transformers 相比,它们的速度更快,内存使用更低。然而,线性注意力表现不佳,而且速度也很慢。尽管如此,我们还是选择追求线性注意力。

一个原因是它的数学吸引力——我们相信它的性能应该更好。另一个原因是,我们认为稀疏注意力的上限是完全注意力,难以超越。另一方面,线性注意力有潜力超越它。

您能解释一下什么是线性注意力吗?

线性注意力本质上是一种核技巧。在 transformers 中,矩阵 Q、K 和 V 的乘法运算,由于维度不同,先乘 QK 还是先乘 KV 会导致不同的计算复杂度。

先乘 KV 可以将计算复杂度降低到线性。然而,问题在于 QK 相乘之后会进行 softmax 运算,这不满足交换律,无法轻易拆分成先乘 KV。因此,线性注意力的第一步是移除 softmax。

但是,移除 softmax 会影响结果。后续的任务是在没有 softmax 的情况下保持结果的一致性,这就是线性注意力旨在实现的目标。

线性注意力、稀疏注意力和线性 RNN 架构之间有什么根本区别?

稀疏注意力本质上仍然是 softmax 注意力。它只是计算的点比密集注意力矩阵少。例如,滑动窗口注意力只计算窗口内的注意力分数,通过减少计算量来实现加速。

线性 RNN 和线性注意力本质上是同一件事,只是有些人称之为 RNN,有些人称之为注意力。

一切都可以写成 RNN 的形式。例如,闪电注意力对应于 RWKV-4,而 RWKV-7 是门控 delta 网络的改进版本。尽管本质上相似,但它们的实现细节有所不同。

线性注意力机制研究的关键里程碑是什么?

大约在 2018-19 年,研究表明可以使用核技巧来降低 transformer softmax 注意力的理论计算复杂度,但结果很差,效率也很低。

在 2019-20 年,稀疏注意力占据主导地位,像 Google 这样的公司提出了许多稀疏注意力变体。后来,线性注意力开始出现,但它面临着性能差和速度慢的挑战。

研究人员主要采用了两种方法进行改进:一种是近似 softmax 函数,使分布符合 softmax;另一种是我们选择的方法,使用完全不同的方法进行建模,而不考虑近似 softmax。

我们在 2021 年 10 月发表了第一篇论文《COSFORMER: RETHINKING SOFTMAX IN ATTENTION》,该论文用余弦函数代替了 softmax 运算,从而可以拆分计算。

在 2022 年上半年,我们发表了第二篇论文《The Devil in Linear Transformer》,该论文分析了线性注意力性能下降的原因,并提供了解决方案。这是闪电注意力的前身。

后来,我们还专门针对线性注意力和长卷积研究了位置编码,发布了 TNN,《TOEPLITZ NEURAL NETWORK FOR SEQUENCE MODELING》,这是一种类似于 S4(Mamba 的前身)的方法。

最后,我们推出了闪电注意力,通过改进的衰减方法和网络结构,使其性能与 transformers 相匹配。我们还使用了一种平铺技术来提高速度。

您如何看待当前非 transformer 架构的技术路线?

线性注意力实际上是一种非 transformer 方法。目前,除了类似 RNN 的方法外,其他非 transformer 架构都在衰落。

例如,像长卷积和大核卷积这样的 CNN,感觉它们由于性能不佳而被逐渐淘汰,但它们实际上在某些方面非常强大,在序列建模中仍然具有一定的作用,例如异常检测任务。

实际上只有三种非 transformer 架构:线性注意力、长卷积和线性 RNN。

但实际上,这三者可以统一为一个,我们称之为线性复杂度模型。我们写了一篇文章涵盖了所有这三者。

闪电注意力与 Mamba 和 RWKV 之间有什么核心区别?

最核心的区别在于,闪电注意力是最简单的线性注意力。Mamba 和 RWKV 都使用数据相关的衰减,而闪电注意力使用手工制作的衰减来提高速度。

虽然可学习的衰减可以实现更好的效果,但它牺牲了速度。例如,RWKV-7 比门控 delta 网络慢 10-15%,而门控 delta 网络的速度约为闪电注意力的一半。

RWKV 的建模效果确实比闪电注意力更好,但它速度较慢,并且尚未解决检索问题。

现在业界是否普遍认为线性注意力具有高且可行的上限?

不,如果业界普遍认同,每个人都会扩大线性注意力模型的规模。而且现在也没有达成共识。如果业界普遍认同,每个人都会做线性注意力,但正如你所看到的,情况并非如此。

但对我们来说,我们已经在 2023 年下半年看到了这一点。当时,我问了很多人,也和很多人交谈过,他们提出的最常见的观点是,他们知道线性注意力在小规模上确实有效,但他们认为一旦扩大规模就会失败。

当时,我想我会扩大规模,让每个人都看到。现在 MiniMax-01 已经问世,没有人怀疑线性注意力在大规模上的能力。

从小型实验到大规模实施

您认为线性注意力的上限可以超过完全注意力吗?

我们现在可以看到,混合架构优于纯 transformers。但纯线性注意力最大的问题是检索能力,这对学术界来说是一个难题。

现有方法虽然复杂且缓慢,但仍然无法完全解决这个问题,这就是为什么有必要转向混合架构。

您观察到什么节点,让您决定走出实验室?

在 2023 年 5-6 月,我们内部已经有了闪电注意力 2,这是世界上第一个比 Flash attention 更快的线性注意力实现。

我们认为它已经跨越了工业红线,其技术成熟度非常高,可以扩大规模。

您如何定义这条工业红线?

首先,效果优于 transformer,其次,速度快于 transformer。这使其具有替代 transformer 的能力。我们当时在一个 15B 规模的密集模型上验证了这一点。

在您走出实验室的节点,为什么最终与 MiniMax 走到了一起?

实际上,当时我已经与一些大公司谈过了。但最终,我还是和 MiniMax 一起完成了这件事。

首先,cosformer 是我与俊杰合作撰写的一篇文章。我们有合作的基础。俊杰在商汤科技时是我的老板。23 年底,俊杰邀请我吃饭。他对这些前沿技术的可能性更有信心。我的理解是,他当时也在寻找技术突破。

当时,MiniMax 已经完成了 Moe 的研究,下一步实际上几乎没有技术突破点。当时,闪电注意力已经发布,mamba 也非常流行,所以在他的眼中,这是一个可行的方向。

这与 MiniMax 的互动伴侣产品有关吗?

没有联系。闫俊杰更关心模型的上限,以及如何进一步突破这个天花板。

在公众眼中,线性注意力可能更多的是突破效率的方向,而不是突破天花板。

这里的重点是,首先,每个制造商的计算能力是恒定的。模型加速得越快,它就能处理更多的数据,产生的模型就越好。当计算能力恒定时,模型速度越快,效果就越好。

您是否观察到数据已经达到峰值的情况?

还没有吧?数据仍然处于持续扩展的阶段,但可能不像 23 年那样激进。

因为数据总是在增加,而且每天都有新数据产生。对于模型来说,它每天都有新数据需要处理。互联网每天产生的数据量非常大。通过清理,我们仍然可以获得新的数据。

与人类发展了这么多年积累的数据相比,数据增长速度是否放缓?

实际上,不一定。看看中国五千年的历史,只积累了那么几本书。但随着互联网的发展,数据量的增长是一条非常陡峭的曲线。互联网之前产生的总体数据可能不如互联网之后一年产生的数据多。

在扩大规模的过程中,闪电注意力面临哪些挑战?

为了验证其可扩展性,我们首先进行了 scaling law 实验,逐渐从小模型扩展到 7B、9B,最后扩展到 400B 以上的模型。

而且我们在理论上证明了线性注意力的容量大于 transformer。

我们将容量定义为 RNN 当前状态的大小。对于 transformer,容量大小为 O(d),其中 d 是大小;对于线性注意力,容量大小为 d²/h。由于 d 远大于 h,因此容量更大。

最后,我们还验证了混合模型优于纯 transformer。

如何实现 4M 长度的序列窗口?

对于闪电注意力,训练长度可以是任意的。只要计算能力得到充分利用,训练 8K、32K 或 128K 的速度是相同的,TGS(每 GPU 每秒处理的 token 数)也是相同的。

因为 transformer 是 n² 计算复杂度,序列越长,计算复杂度增长越快,延迟以二次曲线增长。在 1M 长度下,softmax 注意力的延迟是闪电注意力的 2700 倍。

未来要实现无限上下文窗口,还需要解决哪些技术挑战?

在我们当前的混合架构中,仍然有 1/8 的 softmax 注意力。这在 1M 长度时是一个瓶颈。这 1/8 带来的延迟远高于剩余 7/8 的线性注意力。

如果我们想优化长文本,我们必须考虑优化 softmax 注意力部分。我们可以学习稀疏注意力方法,使其更快更轻。

此外,我们也在考虑使 softmax 和线性注意力的混合比例更加极端,不再是 1/8,而是可能是 1/16 或 1/32。最激进的解决方案是在整个模型中只放置一层 softmax,但为了保险起见,我们没有采用它,主要是考虑对检索能力的影响。

为什么检索能力对模型如此重要?

检索是上下文学习的基础,是必要条件。

你必须记住上下文中的信息才能进行上下文学习,而上下文学习是当前大型模型所有高级能力的基础,例如 CoT(思维链),尤其是长 CoT,它们都依赖于检索能力。

决定性的新架构

您是否关注了业界 FFN 和注意力方面的最新架构改进?

FFN 的改进是 Moe。我也关注了 Byte 的 Ultra Mem,但我认为它是一种有损的东西,一种有损压缩。如果将来扩大规模,可能会出现问题,但我们还没有扩大规模,所以我只能说可能会出现问题。

因为 FFN 基本上就是这些。我们在 Moe 领域的改进无非是从以前的大专家模式转变为当前的小专家模式,使其更加稀疏,然后进行一些加速,这需要进一步研究。

如果你想进一步优化它,因为 FFN 是矩阵乘法,优化只能在 Nvidia 的 CUDA 级别上进行,进行一些矩阵乘法的底层优化。

您是否关注了业界注意力架构方面的改进?

注意力方面的改进基本上是线性的。我们也在考虑将来是否要制作更强大的 Linear,并在当前的基础上进一步加速线性注意力。

有很多改进的方法,一种是改变衰减,另一种是改变一些内部的小技巧。你可以期待我们的新论文。

我们当前上下文长度与推理成本的比率是否相对先进?

一旦涉及到延长序列长度,我们就有非常明显的计算能力成本优势。序列越长,成本优势就越明显,无论是推理还是训练。

例如,在 1M 长度上,线性注意力消耗的计算能力是完全注意力的 1/2700。相比之下,因为我们仍然有 1/8 的完全注意力,所以基本上是 transformer 架构的 1/8,因为线性注意力基本上不计入成本。

如果计算成本如此之低,它可以实现计算瓶颈吗?

现在它确实是一个内存访问瓶颈。解码是一个内存访问瓶颈,而不是计算瓶颈。因为闪电注意力非常快,它太快了,以至于内存访问占用的资源与计算一样少。这主要是因为实际应用中的序列长度不够长。

未来如何使其成为计算瓶颈取决于如何优化内存访问。这些将是工程部门需要负责的事情。

如果线性架构成为下一代的主流架构,哪些硬件适配改进更适合它?

这里有一个非常棘手的事情是,我们需要考虑序列长度。如果你的序列长度集中在 8K 或 32K,那么注意力只占百分之十几,剩下的百分之八十是 FFN 部分。

即使你将注意力优化到极致,优化到 0,你也只优化了百分之十几的延迟。但如果你延长序列长度,注意力的比例会越来越大。这是与完全注意力相比,但对于线性注意力,它的比例是不变的。

因为 FFN 也是线性的,线性注意力也是线性的,它的比例大约是 10%,即使在 1M 的情况下也几乎没有变化。

但如果是完全注意力,注意力计算可能占 99%,下面的 FFN 只占 1%。所以线性注意力只有在长文本中才有优势。

如果线性架构成为主流,那么追求的可能是低能耗硬件,只降低能耗。包括 Spiking Neural Network (SNN) 芯片可能更适合,实际上有些人正在做这件事。

展望通往 AGI 的道路

您对模型开源效果有什么期望?

首先是宣传效果。我个人认为,除了展示一些实力外,开源最重要的是看将来大家如何使用它。我认为小型模型开源可能是我们将来更考虑做的事情。

以及如何为每个人 finetune 制作一些基础设施也可能需要考虑。开源对我们来说是未来的一件长期的事情,旗舰模型应该继续开源。

将来是否有可能出现非混合的纯血架构?

目前,没有任何方法可以比混合架构做得更好,尤其是在速度方面。添加一小部分 softmax 注意力,当序列长度不是特别长时,速度优势非常明显,尤其是在 flash attention 出现之后。

对纯血架构的研究仍在进行中,但这非常困难,而且没有更多的唾手可得的果实。我们有一些技术解决方案,但实现起来并不简单,最终取决于我们需要实现多长的序列长度。

另一个问题是,对超长文本是否有强烈的需求?尽管像 Claude 这样的模型已经达到了 200K 上下文,但用户似乎对当前的长度非常满意。Agent 应用将来可能会带来对超长序列的需求,但还没有成熟的基准。

但我认为这个问题就像 Nvidia 为未来的游戏开发先进性能的显卡一样,即使现在不需要,它也是未来的技术。

例如,深度研究需要模型阅读几十个网站的内容,处理时间在几十分钟的量级,这可能是长文本的一个应用方向。

您认为 CoT 之后下一个重要的事情可能是什么?

我们已经考虑过这个问题。首先,当前的推理模型相对流行,今年主流仍然是推理部分。在那之后,我们很难想到纯语言模型在未来会有任何特别大的变化。

我也与其他老师交谈过,他们的感觉是,每个人都会重新降低模型的成本,以便推理速度变得越来越快,价格变得越来越低,并在降低成本的同时保持效果。

因为天花板正在迅速逼近,绝大多数情况是在检查和填补大型模型的能力空白。但如果出现更大的技术突破,短期内可能会比较罕见,我们还没有看到它们。

在 MiniMax 探索了线性注意力之后,下一个探索方向可能是什么?

接下来可能是探索多模态的架构,特别是我们是否想要做这种原生生成和理解统一的大型模型架构。

以 AGI 为终点,计算复杂度为 O(n²) 或 O(n) 的模型哪个会是更好的答案?

当然是 O(n)。从拟人化的角度来看,人一定是 O(n) 复杂度。例如,如果一个人的复杂度是 O(n²),那么我与你说话的速度就会变得越来越慢。

因为对于 transformer 来说,其推理复杂度是 O(n²) 计算复杂度,也就是说,吐出第一个 token 和吐出第 100 个 token 的延迟是不同的。

我们人类无法想象这样的事情,因为人们从出生以来就没有重启过,并且一直在吐东西,所以人们的计算复杂度是恒定的。

人一定是智能的最佳解决方案吗?

我们目前只能这样认为。也有一些人正在走仿生智能的路线,但我们并没有过多关注这些方向。

以 AGI 为终局,模型改进的哪些领域是最重要的事情?

除了语言建模之外,还有学习方法的问题。你如何学习,以及从环境中学习,从与环境的互动中学习非常重要。毕竟,当前的多模态理解仍然非常缺乏数据。

而且即使是机器的 few-shot 学习目前也是被标记的,但人类的学习是无标记的。因此,如何将一切统一在一个自我构建的框架下也是一个问题。