Transformer 的上限在哪里?这个问题从 2022 年 ChatGPT 问世起,就开始被广泛讨论。
最近一两个月,这个问题的答案初见端倪。前 OpenAI 推理负责人 Jerry Tworek 与前 Google Gemini 预训练负责人 Rohan Anil 公开表态,Transformer 已经走到尽头,并成为走向 AGI 的最大瓶颈[1]。
国内头部基础大模型也在积极尝试改进版的 Transformer 架构,但进行这些改进的一部分原因是算力瓶颈,并不完全是为了提升架构的能力上限。通过稀疏/线性注意力来降低架构复杂度,这些方案大幅提升了训推效率,在资源受限的情况下成功实现了更大的参数规模,并让模型能力逐渐趋近国外闭源模型。然而,半年后,这些复杂度更低的架构是否同样会碰壁,还需要打一个问号。
那么,随着 Transformer 上限触顶,AI 的下一代模型架构应该具备哪些特质?来自上海人工智能实验室的书生 Mobius 团队给出了一条可能的路径——通过知识与推理分离,探索在持续学习、组合泛化等方向上能力上限更高的架构,并顺带用更高的能力上限倒逼训推效率的提升。
在这个思路下,团队提出的 Mobius 系列架构,有望在未来赋能商用落地、持续自进化、科学发现等多个关键问题或领域。
自从 Transformer 问世以来,工业界通过不断 Scaling 数据量、参数量,以一种 ROI 非常高的方式,获得了举世瞩目的强大模型。这些模型在自然语言、代码、数学、视频等许多任务上,展现出了前所未有的超强智能。
然而,在基于 Transformer 的模型越来越大、能力越来越强的同时,Transformer 始终具有的「三宗罪」却一直被忽略,没有得到有效解决。这「三宗罪」分别是:
(1)学过的,不能快速回想起来;
(2)新学的,想要记好,成本太高;
(3)没学的,难以通过联想构建。
对于(1),从现有模型的端到端推理效率中便可见一斑。初代 ChatGPT 仍然不具备很强的数学和代码能力。随着 OpenAI O1 的横空出世,超长思维链(Chain-of-Thought,CoT)以及配套的强化学习(Reinforcement Learning,RL)算法正式进入公众视野。
在深度思考模式下,通过不断输出 Token,模型可以持续进行试错、验证和纠错,并大幅提高答对复杂问题的概率。可见,模型并不是不会解决这些问题,而是需要恰当的手段激发出它们内在的潜力。
然而,极致的「啰唆」意味着极致的不聪明和反应慢,长 CoT 只是一种较为低效的智能激发方式。抛开聪明与否不谈,这种超长 CoT 还会让模型在完成任务时耗费非常多的算力和时间。因为长度就摆在那里,推理效率很难随着对 Attention 或 Infra 的充分优化而产生质的飞跃。
对于(2),这其实是在 AI 领域存在已久的持续学习问题,狭义上指在学会新知识的同时不遗忘旧知识,并非 Transformer 自己的「锅」。反而在 Transformer 架构被不断 Scaling 之后,人们发现这个架构原生具备 Few-shot Learning 和 In-Context Learning 能力。
因此,在回答模型没学过的问题时,只要通过 RAG 的方式构造足够充分的上下文,模型就有不小的机会答对这个问题。
然而,人类的需求是无限的,Transformer 的注意力是有限的。随着解决问题所需的知识不断增多,RAG 得到的长上下文所具有的低信息密度属性,反而会成为模型解答问题的瓶颈(即使是开卷考试,不画重点地扔给我们几百页从未读过的书,恐怕也很难找到答案)。看起来,高密度的记忆才能带来真正高效的持续学习。
不过,对于 Transformer 而言,最有效的高密度记忆方式是大规模预训练。假如希望加入一个新领域的知识,最好将核心的旧数据和部分新数据混合,再把模型续训甚至重训一遍。一方面,成本不可接受;另一方面,这也违背了持续学习的初衷。
对于(3),这是一个困扰 AI 领域更久的组合泛化问题,狭义上指通过见过的知识,自动学会没见过的知识,也不是 Transformer 所特有的。并且,随着 Transformer 的 Scaling,人们发现这类模型还可以 Zero-shot 地完成一些从未见过的任务。
不过,目前 Transformer 可以 Zero-shot 解决的问题,更多是流程性的、有明确逻辑链条的问题。在一些强调直觉、联想的领域,如科学发现,Transformer 产生重要泛化的例子屈指可数。
爱因斯坦发现「广义相对论」,不仅因为他同时了解「狭义相对论」和「黎曼几何」这两条知识,还因为这两条知识在某一时刻同时被激活,俗称「灵感迸发」。Scaling 后的 Transformer 架构,已经足以让两条知识同时出现在一个模型中。但在单次推理中,如何提高不同知识同时被看见、同时被激活的概率,是增加不同知识之间产生组合泛化机会的关键。
先给出结论,上面的「三宗罪」可以共同归因于一点——分层的记忆与推理。
根据学术界对于 Transformer 的主流理解,全连接层(FFN)负责记忆存储,注意力层(Attn)负责组合推理。不同层的 FFN 存储不同的记忆,相同的 Token 在不同层可以激活出不同的知识;不同层的 Attn 完成不同形式的组合,相同的知识输入在不同层可以产生不同的推理结果。
并且,由于 Transformer 的层间存在递归特性,当前层的 Attn 主要处理上一层 FFN 中取出的知识。当然,由于层间 Residual Connection 的存在,Attn 也可以处理一部分从靠前层取出的知识。
根据 Transformer 的层次化记忆与推理机制,我们可以发现:
(1)在推理时,如果想要完整回顾学过的所有知识,需要经过所有层才行。一旦错过关键知识,只能强行「挤」出一句废话来争取时间,再重新来到第一层,多取几次知识;
(2)知识在不同层的存储管理比较混乱,既存在一定的重叠,又存在一定的耦合关系。学习新知识时,模型无法判断应该将其插入哪里,所以只能暴力地同时更新所有知识;
(3)不同层之间的知识和推理,主要通过 Residual Connection 这条狭窄通路进行连接。即使关键知识都被激活了,等这些知识「见面」时,靠前层的知识也已经被稀释得差不多了。
可见,分层的记忆与推理,不仅限制了知识读取的效率,还限制了知识存储的便利性,更限制了不同知识之间产生组合泛化的概率。如何构建更好的记忆和推理机制,可能成为提升 AI 模型能力上限的关键。
假如将不同层中记忆与推理的耦合关系拆解开,将每层的记忆全部取出,存储为一个共享池,并让每一层的 Attn 都可以访问这些知识,上述三个问题就有机会迎刃而解:
(1)假设推理层数相同,共享记忆后的模型可以更多次地遍历(Traversal)这些知识,有更大的机会取出关键知识,也就更容易输出有用的话;
(2)所有知识进行扁平化存储后,新知识就可以「一览无余」地看到自己应该去哪个位置,更少担心自己的加入会产生牵一发而动全身的问题;
(3)在每次经过共享记忆库时,所有知识都有机会被同时激活,大大增加了不同知识直接「见面」的机会,也有可能带来更好的组合泛化。
架构示意图
在这个思路之下,Mobius 架构被设计出来。目前的实验已经证明,Mobius 对(1)有直接帮助:Mobius 在端到端推理效率较 Transformer 提升近 4 倍的同时,下游任务准确率并未变差。
并且,Mobius 对(2)和(3)也已经展现出间接帮助。例如,Mobius 可以用 60% 的数据达到与 Transformer 相近的 MMLU 分数,具备更好的数据压缩率;在一些知识组合泛化任务中,Mobius 相较 Transformer 提升了 2 倍。
7B 从头预训练的 MMLU 得分
35B 续训练的下游任务得分
35B 续训练的推理速度
35B 续训练的 CoT 长度
Toy Model 的组合泛化验证
其实,从 Transformer 到 Mobius 的改造,与当初从 RNN 到 Transformer 的改造有着异曲同工之处:
(A)从 RNN 到 Transformer 的改造,是将 Token 之间「Recurrent + Residual Connection 的组合」转换为 Full-Connection;
(B)从 Transformer 到 Mobius 的改造,是将知识之间「Recurrent + Residual Connection 的组合」转换为 Full-Connection。
上一代改造完成后,Transformer 大幅提升了基座模型的能力上限。这一次改造会不会产生相似的效果,我们拭目以待。
RNN、Transformer、Mobius 的对比
近期,对于 Residual Connection 和 Latent Reasoning 的改进比较火。如果站在这两个角度细想,你可能还会发现,Mobius 原生支持 Backward Residual Connection 和 Dynamic Latent Reasoning。
无论是最早的 ResNet,还是近期出现的 Hyper-Connection、Attention-Residual,它们的共性功能都是「在前向传播时,将浅层的信息传向深层;在反向传播时则相反」。
正如上文分析过的,这种单向传播存在一个问题:一旦有些关键知识被错过,或者有些深层知识反而是浅层知识的读取前提,模型就只能输出一些「废话」,再循环一次。
Mobius 在做法上,是将不同层的知识存储到一起;但在思想上,可以认为它引入了一种反向的 Residual Connection。这种残差连接可以让信息传递更加灵活,提高模型的知识压缩率和推理效率。
相比 Token-by-Token 的推理方式,Latent Reasoning 使用连续的 Embedding 连接两次循环,也起到了某种残差连接的作用。不过,传统的 Latent Reasoning 是低效的。
首先,Latent Reasoning 要求模型每次循环都经过所有层,不论某些层的知识是否已经不再必要。另外,进行多次迭代的 Latent Reasoning,通常只是为了迭代一个 Token 的表征。相比之下,Mobius 原生支持一种更加动态的 Latent Reasoning 机制。
一方面,由于 Mobius 在每一层都有机会访问所有知识,并且没有限制这一层必须激活哪些知识,因此,Mobius 每经过一层,不仅可以视作完成了一次 Latent Reasoning,而且还能进行更加灵活的激活选择,减少冗余计算。
另一方面,Mobius 在每一层迭代的不只是一个 Token 的表征,因此,在经过各层的过程中,它具备更高的迭代效率和更加连续可微的优化路径。
递归自进化、AI 辅助科学发现和世界模型,正成为最重要、也最亟待解决的前沿应用:
不难发现,Mobius 知识与推理分离的设计,恰好呼应了递归自进化和科学发现的两个核心需求,在这两个方向上有机会成为比 Transformer 更具潜力的模型。Mobius 原生支持的 Latent Reasoning,则有机会为世界模型带来帮助。
不过,如果想在更大程度上帮助世界模型,对 Attn 机制的改变可能更加重要,这些变化也会在未来版本的 Mobius 中有所体现。
现阶段的 Mobius 虽然带来了端到端推理速度的提升,但这部分提升主要源于 Mobius 更加高效的推理路径带来了更短的 CoT。在比较单个 Token 的推理效率时,Mobius 相比 Transformer 还有一点距离。
这部分距离主要源于 Mobius 在算法上的知识与推理分离,天然带来了硬件上的存算分离特性。这种特性所面临的访存墙、通信墙,是阻碍 Mobius 进一步提高推理效率的根源。
Mobius 提供了更高的能力上限,但如何设计更加适合 Mobius 的存取策略、并行策略,甚至更加适合 Mobius 的硬件架构,是决定 Mobius 能否达到自身能力上限的关键。
[1] https://www.youtube.com/watch?v=2RJiaf0SY8s