Mamba 架构实现推理性能超 Gemma3-27B!推理模型开始迈入「无注意力」时代

B站影视 港台电影 2025-06-07 12:58 2

摘要:推理大模型中的注意力机制可以去掉吗?这是当下最具挑战性的开放问题之一。过去半年,伴随长思维链的爆火,模型在复杂推理任务中普遍需要生成成千上万的 tokens 才能产出高质量解答。然而,这也与 Transformer 架构天然矛盾:随着生成长度的增长,Atten

PromptCoT-Mamba是首个实现解码显存常量、计算复杂度线性、长上下文稳定扩展、且具有强推理能力的模型。

推理大模型中的注意力机制可以去掉吗?这是当下最具挑战性的开放问题之一。过去半年,伴随长思维链的爆火,模型在复杂推理任务中普遍需要生成成千上万的 tokens 才能产出高质量解答。然而,这也与 Transformer 架构天然矛盾:随着生成长度的增长,Attention 模块需要不断堆积 Key-Value Cache,带来线性增长的显存消耗与二次方增长的计算复杂度,直接限制了推理深度、解码速度与部署成本。

半年以来,学术界与产业界在持续探索“去注意力机制”的模型架构:腾讯混元提出 Hunyuan Turbo-S、Nvidia 发布 Nemotron-H、Tri Dao 也在 M1 中尝试了极致的 Mamba-Attention 混合架构。这些混合模型虽然显著减少了 Attention 占比,但仍未能完全摆脱对注意力机制的依赖,Key-Value Cache 依然存在,解码复杂度仍受限于推理链长度。

与此同时,Diffusion LLM(dLLM)路线也在最近快速升温,试图跳脱传统自回归范式,通过 Masked Denoising 等迭代去噪机制优化解码过程。然而,本质上大多数 dLLM 仍构建在 Transformer 架构之上,核心依赖注意力机制完成每轮去噪迭代。在解码阶段,dLLM 需要多轮完整的序列更新,每轮都需重新激活 Attention 计算,导致计算复杂度依然呈二次增长、Key-Value Cache 持续膨胀,长推理链生成的内存开销和解码时延甚至较自回归模型更高。换言之,现有 dLLM 尚未摆脱注意力机制的束缚。

那么,究竟有没有可能完全移除注意力机制,同时在复杂推理任务上做到高效与强大兼备?最近,蚂蚁技术研究院联合香港大学发布了 PromptCoT-Mamba-7B,首次实现了解码显存常量、计算复杂度线性、长上下文稳定扩展的推理大模型。至此,推理大模型终于可以完全不依赖注意力机制而运作起来了!

1 模型概览

无注意力架构:全模型仅基于 Mamba-2 的 SSD 层,无需 Key-Value Cache,彻底消除随生成长度增长的内存负担;

线性时间复杂度:得益于状态递推设计,每生成一个token的计算复杂度为 O(NP),而非 Transformer 的O(TN),其中 T 为序列长度,N,P 为隐藏维度。解码速度在长序列条件下具备绝对优势;

强大的推理能力:首次在竞赛数学与代码推理任务上全面超越同尺寸及更大规模的 Transformer 以及混合架构;

开放通用的训练框架:将 PromptCoT 扩展为跨领域数据合成范式,不仅支持奥数题,更系统性地扩展到了复杂代码生成任务,支撑起完整通用的推理模型训练框架。

技术报告地址:https://arxiv.org/abs/2505.22425

GitHub:https://github.com/inclusionAI/PromptCoT

2 模型效果

1. 主实验:首个击败Transformer的无注意力推理大模型

在主实验设定下,PromptCoT-Mamba-7B 在无需任何注意力机制与 Key-Value Cache 的前提下,首次在竞赛级数学与代码推理任务中全面超越同尺寸甚至更大尺寸的 Transformer 系列模型,验证了注意力之外的结构化建模潜力:

在4个数学推理评测集 MATH-500,AIME 2024,AIME 2025,以及 OlympiadBench,以及3个代码推理评测集 LiveCodeBench-v5, HumanEval, 以及 HumanEval+ 上全面超过 s1.1-7B 以及Nemotron-H-8B等 Transformer 和混合结构。

和 Transformer 强基线模型 s1.1-7B 相比,AIME 2024 提升16%, AIME 2025 提升7.1%,LiveCodeBench-v5 提升16.6%。

在 AIME 2024, AIME 2025,以及 LiveCodeBench-v5 上甚至超过了 Google 的 Gemma3-27B。

2. 数学专项提升:极限推理能力释放

在进一步领域特化设定下,数学专属版本 PromptCoT-Mamba-Math-7B 显著提升模型在数学复杂推理任务中的表现:

AIME 24:提升至 42.9%(相比通用版提升 +7.7%)

AIME 25:提升至 30.8%(相比通用版提升 +6.2%)

3. 常量内存推理效率:无注意力架构优势凸显

得益于完全去除注意力机制及 Key-Value Cache,PromptCoT-Mamba-7B 在长序列解码与低内存部署环境下展现出显著的效率优势:

在 24GB GPU 低资源设定下,相较 Transformer,推理吞吐达到其3.66倍

在 72GB GPU 设定下,相较 Transformer,推理吞吐达到其1.69倍

这使得 PromptCoT-Mamba 成为长上下文与实际工程部署的高效架构选择。

3 更多技术细节

1. PromptCoT问题合成框架

PromptCoT [1] 是研究团队于2025年3月推出的问题合成框架。其基本思想是从概念出发,通过生成 Rationale 来模拟人类专家在命题过程中的抽象思路与逻辑组织,最终 Rationale 类似于设计架构,和输入概念一起指导竞赛级数学问题的合成。

在理论上,PromptCoT 定义了最优 Rationale 的生成目标,要求其同时最大化:

Rationale 在给定概念条件下的生成概率;

完整问题在已给定 Rationale 与概念条件下的生成概率。

在实现上,首先通过大语言模型生成针对现有竞赛级题目的 Rationale 样本,随后利用这些概念-Rationale-问题三元组训练问题生成模型,使其具备从任意概念集合自动构建高难度推理问题的能力。该机制不仅提升了样本生成的逻辑性与难度控制能力,也大幅缓解了复杂推理训练过程中的数据瓶颈问题。

PromptCoT 的第一版在 DeepSeek-R1-Distill-Qwen-1.5B, DeepSeek-R1-Distill-Qwen-7B,以及 QwQ-32B 基础上,仅通过 SFT 就实现了数学推理效果上的飞跃。PromptCoT-DS-7B 在 MATH-500,AIME 2024,以及 AIME 2025 上的表现甚至超过了32B的 s1 模型。

在此基础上,经过版本迭代,PromptCoT 实现了无注意力模型架构上数学和代码推理能力的同时突破。

2. 两阶段训练流程

PromptCoT-Mamba 训练分两阶段:

初始阶段:以多领域覆盖式的高质量推理轨迹数据进行基础训练,约 188 万样本,涵盖数学、代码、逻辑、复杂文字推理任务,帮助模型快速掌握通用推理范式。

进阶阶段:引入 PromptCoT 合成的 25 万高难度问题样本,进一步雕琢复杂长链路推理能力。

推理阶段使用最大长度达 64K tokens 的长序列推理,配合 SSD 架构特性,稳定高效。

4 结语展望:为无注意力推理大模型生态奠定技术基础

PromptCoT-Mamba 不仅证明了无注意力架构在复杂推理任务上具备工业级应用潜力,更为未来 Mamba 体系的推理技术研究提供了完整可复用的技术栈与数据生成范式:

架构端:无注意力推理的可行性范式;

数据端:PromptCoT 跨领域高难任务生成框架的普适扩展。

未来,任何复杂推理任务只需通过 PromptCoT 机制扩充领域概念库,即可快速构建新型纯 Mamba Reasoning 模型,成为无注意力推理大模型的新一代基座方案。

5 作者介绍

该工作第一贡献者为香港大学计算机系博士生赵学亮,师从孔令鹏教授;蚂蚁技术研究院武威为共同贡献者。

[1] Xueliang Zhao, Wei Wu, Jian Guan, and Lingpeng Kong. PromptCoT: Synthesizing Olympiad-level Problems for Mathematical Reasoning in Large Language Models. To Appear in ACL 2025.

未经「AI科技评论」授权,严禁以任何方式在网页、论坛、社区进行转载!

公众号转载请先在「AI科技评论」后台留言取得授权,转载时需标注来源并插入本公众号名片。

//

来源:AI科技评论一点号

相关推荐