本文提出 PROMPTCOT-MAMBA，一种基于 Mamba-2 状态空间模型的无注意力语言模型，通过两阶段课程微调和 PROMPTCOT 合成范式，在数学和代码推理任务上超越同规模甚至更大规模的 Transformer 模型，同时实现固定内存和高效推理。
State Space Model, Reasoning, Efficiency, Fine-tuning, Generative AI
Xueliang Zhao, Wei Wu, Lingpeng Kong
The University of Hong Kong, Ant Group
Generated by grok-3
Background Problem
大型语言模型（LLMs）在复杂推理任务（如数学问题求解和代码生成）中取得了显著进展，但 Transformer 架构的注意力机制带来了两个核心瓶颈：内存密集且随上下文长度线性增长的键值（KV）缓存，以及在高难度领域缺乏结构化的微调策略。这些问题在长上下文链式推理（Chain-of-Thought）任务中尤为突出。本文提出了一种无注意力机制的语言模型 PROMPTCOT-MAMBA，旨在通过状态空间模型（Mamba-2）实现固定内存和常数时间推理，同时通过课程微调提升复杂推理能力，解决上述效率和性能问题。
Method
PROMPTCOT-MAMBA 的方法主要分为两个核心部分：
-
架构创新：基于 Mamba-2 的状态空间双层（SSD）
该模型完全摒弃了自注意力机制，采用 Mamba-2 的 SSD 层作为计算单元。SSD 层通过递归状态更新（见公式 (1)）实现自回归推理，具体步骤为：输入词嵌入通过线性投影生成时间依赖参数（如衰减因子 和核向量 ），然后通过低秩状态更新和线性读出计算下一 token 的分布。其复杂度为 ，内存占用固定，推理时间与序列长度无关，相比 Transformer 的 复杂度和线性增长的 KV 缓存具有理论上的效率优势。训练时，SSD 层通过结构化收缩操作实现并行计算，复杂度为 ，仍优于 Transformer 的 。 -
课程微调：基于 PROMPTCOT 合成范式的两阶段策略
训练分为初始化阶段和高级阶段。初始化阶段使用开源数据集（如 OpenCodeReasoning 和 OpenThoughts2）进行基础推理能力微调，包含大量自动合成的示例。高级阶段引入 PROMPTCOT 合成范式，通过两步生成过程（先基于抽象概念生成专家级推理过程，再条件化生成目标问题）构建高复杂度的训练数据，旨在培养专家级推理能力。PROMPTCOT 的数据生成强调教学结构和抽象深度，但其质量和潜在偏差未在论文中充分讨论。
批判性思考：虽然 SSD 层在效率上有优势，但其在捕捉长距离依赖和复杂上下文关系方面的能力可能不如注意力机制，尤其是在需要全局信息的推理任务中。课程微调策略依赖于合成数据的质量，若 PROMPTCOT 生成的问题缺乏多样性或存在模式偏差，可能导致模型过拟合而非真正提升推理能力。此外，论文未讨论 SSD 层在训练稳定性或收敛速度上的潜在问题，这可能是实际应用中的隐患。
Experiment
实验在七个数学和代码生成基准测试（如 MATH-500、AIME 24/25、Livecodebench）上评估了 PROMPTCOT-MAMBA-7B 的性能，采用 pass@1 准确率作为主要指标，并对部分基准使用 avg@k 评估以减少生成随机性。实验设置包括与多个 Transformer 和混合 Mamba-Transformer 基线（如 Gemma3-27B、S1.1-7B）的对比，训练在 8×A100 80GB GPU 上使用 DeepSpeed ZeRO Stage 2 进行分布式优化。
结果：PROMPTCOT-MAMBA-7B 在多个高难度任务上表现优异，如 AIME 24（35.2%）、AIME 25（24.6%）和 Livecodebench（29.9%），超越了同等规模的 Transformer 模型，甚至超过更大的 Gemma3-27B（分别高出 2.6%、0.6% 和 3.0%）。消融研究表明 PROMPTCOT 合成阶段对性能至关重要，去除后准确率大幅下降（如 AIME 24 从 35.2% 降至 11.7%）。效率测试显示其推理吞吐量在 24GB 内存下比 S1.1-7B 高 3.66 倍，在 72GB 下高 1.69 倍。
分析与批判：实验结果在特定任务上令人印象深刻，但基准选择偏向数学和代码推理，缺乏对自然语言理解或更广泛任务的评估，可能掩盖了模型在其他领域的局限性。消融研究虽证明了 PROMPTCOT 的重要性，但未探讨合成数据可能引入的偏差或过拟合风险。此外，效率测试虽显示优势，但未考虑实际部署中的复杂因素（如批处理优化、硬件适配），可能高估了实际收益。实验设计总体合理，但不够全面，需更多样化的任务和更深入的鲁棒性分析来验证模型的泛化能力。
Further Thoughts
PROMPTCOT-MAMBA 的研究为探索注意力机制的替代方案提供了一个有价值的视角，尤其是在长上下文推理任务中的效率提升方面。然而，其依赖于状态空间模型的架构可能在某些需要强全局依赖的任务（如长篇文本理解或多模态推理）中遇到瓶颈。未来的研究可以探索如何结合少量注意力机制或设计混合架构，以在效率和性能之间取得更好平衡。此外，PROMPTCOT 合成范式的数据生成策略值得进一步优化，例如引入多样性约束或对抗性生成机制，以避免潜在的模式偏见。
另一个有趣的方向是将此方法与其他领域（如多模态系统或机器人控制）结合，测试状态空间模型在处理时间序列数据或实时交互任务中的潜力。例如，结合近期在多模态基础模型中的进展，是否可以将 PROMPTCOT-MAMBA 扩展到视觉-语言推理任务中？同时，与联邦学习或隐私保护机器学习结合，探索其在资源受限设备上的高效推理能力，也是一个值得深入的方向。总之，本文开启了无注意力模型在复杂推理中的应用讨论，但其局限性和适用范围仍需更多研究来明确。