还是原装Transformer好！北大清华团队同时揭示Mamba等推理短板

首页 » 还是原装Transformer好！北大清华团队同时揭示Mamba等推理短板

文章来源于互联网:还是原装Transformer好！北大清华团队同时揭示Mamba等推理短板

AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

北京大学的工作来自智能学院王立威、贺笛老师课题组，发表于ICML 2024。作者包括北京大学图灵班本科生杨铠；苏黎世联邦理工学院硕士生Jan Ackermann；北京大学智能学院博士生何震宇、冯古豪、张博航；纽约大学博士生冯韫禛；北京智源研究院研究员叶启威；清华大学的工作来自于明年即将入职清华大学交叉信息院做助理教授、目前在加州大学伯克利分校 Simons Institute 做博士后的吕凯风研究员。作者包括斯坦福在读一年级博士温凯越；清华大学姚班本科生党星宇。

思维链（CoT）是大模型中最神秘的现象之一，尤其在数学任务上显著提升了 Transformer 的能力。然而，思维链的引入也使生成内容的长度增加，消耗了更多的计算资源。这不禁让人好奇：最新推出的高效模型（如 Mamba）是否也能像 Transformer 一样具备强大的推理能力？近期，北大和清华的研究团队同时给出了明确的否定答案，揭示了 Mamba 等高效模型在结构上的局限性。

论文 1：Do Efficient Transformers Really Save Computation? （发表于 ICML 2024）
论文链接：https://arxiv.org/abs/2402.13934

论文 2：RNNs are not Transformers (Yet): The Key Bottleneck on In-context Retrieval
论文链接：https://arxiv.org/abs/2402.18510

基于 Transformer 的大语言模型在文本翻译、文本生成等许多领域展现了惊人的能力。主流的大语言模型通常采用自回归范式进行生成：由问题描述、相关提示组成的输入序列（prompt）会被首先编码。基于编码后的信息，大模型逐步生成后续的单词序列，以形成问题的答案。对于复杂的问题，已有的实践和理论研究表明，利用思维链提示（CoT）可以显著提升模型在数学或推理方面的问题解决能力。

如上图，当给出的问题示例中只有最终结果，而没有中间的推理过程时，大模型在见到新问题时往往会生成错误的答案；但如果给出的问题示例中包含了完整的思维链（中间推导步骤），大模型在新问题上也往往会输出完整的思维链并得到正确答案。然而，思维链提示依赖于大量中间步骤的输出，使用 Transformer 架构会产生大量的计算开销。因此，一个自然的问题便是：能否使用更少的计算量，实现思维链推理？许多架构致力于降低注意力机制带来的计算复杂度，能否在这些架构上使用 CoT 提升性能并实现推理的加速？

近期，来自北大和清华的研究团队从理论角度对上述问题进行了深入探讨。结果令人惊讶：两个团队一致证实包括 Sparse Transformer、Linear Transformer、Mamba 在内的许多架构，即使在这些模型上应用思维链，其理论上的能力上限仍无法解决多种实际推理问题，并与标准 Transformer 有本质差距。这些理论结果为高效结构的实用价值蒙上了一层阴影。

Transformer + CoT 依然是最佳选项

北大的研究团队将推理任务通用地建模为动态规划。在推理过程中，模型需要按照合法的拓扑排序逐步输出子问题的结果（即思维链），以最终得出原问题的解。他们从理论角度证明，若希望 Sparse Transformers、Linear Transformers 等模型通过思维链来解决动态规划问题，模型的宽度必须增加，进而导致时间复杂度恰好达到平方级别。这一时间复杂度与标准 Transformers 在相同长度的推理任务中所需的复杂度一致，表明这些所谓的高效结构在一般推理问题上并不具备计算优势。

殊途同归，清华的研究团队考虑了 RNN 模型和 Transformer 模型在检索、关联回忆、计数，以及判断一张图是否为树等基本问题上的表达能力差异。他们从理论角度证明了，对于问题规模 n，任意 o (n) 大小的 RNN 模型均无法完成上述任务，即便使用任意长的 CoT。但一个固定大小的 Transformer 可以不使用 CoT 解决检索、关联回忆、计数等问题，并使用 O (n) 长度的 CoT 正确判断一张图是否为树。这些结果表明类 RNN 的结构使用思维链获得的能力，距离 Transformer 仍有巨大差距。

两个研究得出相似结论的背后，揭示了相似的本质。这些模型架构之所以具有内存高效性，是因为它们能够在较少的空间下完全确定输出序列。然而，输出序列的可能性指数级增长，这使得这些架构无法以较小的模型尺寸正确生成所有可能的输出。换句话说，模型的规模必须随着问题规模的增加而扩展。这一核心观察表明，包括 Mamba 在内的具有循环（recurrent）性质的网络架构，都受到了相同的限制。

如何提升其它模型使用思维链的能力？

在得到负面结论的同时，两个团队考虑不同角度设计解决方案。

北大研究团队从推理任务的局部性入手，即当前输出所需要向前依赖的最远输出的距离。局部性反映了推理所需要的长程记忆难度。他们从理论角度证明了，当推理任务有较好的局部性时，许多高效模型能够以理论更优的推理速度完美解决问题。

来自清华的研究团队则从引入上下文检索器的角度入手。他们从理论角度证明了，使用显式上下文检索器或使用一层 Transformer 作为隐式上下文检索器，均可以大大增强 RNN 使用思维链后获得的能力。

研究团队还设计了大量的实验验证理论结果。来自北大的研究团队在多种 Transformer 架构上使用 CoT 数据进行训练，表明标准 Transformer 架构事实上使用了最小的计算量。