Blog Tags GitHub
Back to blog
mixture-of-experts llm-inference reinforcement-learning

MAI-Thinking-1: Building a Hill-Climbing Machine

微软AI从零构建的35B/1T MoE推理模型，AIME 2025达97%，SWE-Bench Pro达52.8%

MAI-Thinking-1: Building a Hill-Climbing Machine

一、论文概述

项目内容
标题MAI-Thinking-1: Building a Hill-Climbing Machine
作者The Microsoft AI Team
机构Microsoft AI (MAI)
论文microsoft.ai/pdf/mai-thinking-1.pdf
发布2026年
页数108页
联系mai-technical-report@microsoft.com

二、核心思想

问题定义

如何从零构建一个强大的推理模型，使其在 STEM 推理和编程任务上达到前沿水平？关键不是单个模型的突破，而是建立一个能够持续改进的”爬山机器”（hill-climbing machine）——将模型开发视为系统级优化问题。

解决方案概述

MAI-Thinking-1 是微软 AI 从零开发的推理模型，具有以下核心特征：

  • 架构: 35B 活跃参数 / 1T 总参数的稀疏 MoE 模型
  • 预训练: 30T tokens，完全使用干净的企业级数据，无第三方模型蒸馏
  • RL 训练: 从零开始的强化学习，持续数千步的对数线性性能提升
  • 三大设计原则:
    1. 能力应通过学习获得，而非继承（无蒸馏）
    2. 简单性是可持续的
    3. 科学严谨性避免捷径

核心性能

基准得分
AIME 202597.0%
AIME 202694.5%
LiveCodeBench v687.7%
SWE-Bench Pro52.8%
SWE-Bench Verified73.5%
GPQA Diamond84.2%
HMMT Feb 202684.9%
Terminal-Bench 2.046.0%

三、技术架构

模型架构

架构图

Figure 2: MAI-Base-1 架构总览

MAI-Base-1 是一个 decoder-only Transformer，具有以下关键设计：

参数值
活跃参数34.7B
总参数962B（~1T）
层数78
隐藏维度6656
FFN 维度13312
Expert FFN 维度10240
Top-k / 专家数8 / 512
KV/Q 头数8/80
词表大小200,019 (o200k_base)
最大上下文256K

关键架构创新

1. 交替 MoE + Dense FFN

  • 交替使用高稀疏度 MoE 层和小型 Dense FFN
  • 比均衡稀疏度分配更高效（等活跃参数和等总参数设置下）
  • 每层 MoE 依赖共享专家较少

2. 交替 Local/Global Attention

  • 遵循 Gemma 3 的周期性注意力设计
  • 5 个 Local Attention 层 + 1 个 Global Attention 层
  • Local: RoPE，滑动窗口 512，基频 10,000
  • Global: 无位置编码（效率更高）
  • GQA: 8 KV heads，每头维度 128

3. LatentMoE 设计

  • 共享 down-projection 在 all-to-all dispatch 前应用
  • 压缩因子 2x，专家内扩展 3x
  • 路由决策基于原始表示
  • 8/512 experts 激活，softmax gating

4. Dropless MoE

  • 完全无 token dropping
  • 可变消息大小的 all-to-all 通信
  • 全局 batch 负载均衡损失

模型家族配置

模型活跃总参数层数隐藏Top-k/Experts
L12365M3.9B1210248/512
L241.5B30B2420488/512
L426.1B159B4235848/512
L6621.7B615B6656328/512
L7835.6B1015B7866568/512
MAI-Base-134.7B962B7866568/512

注意力初始化创新

注意力初始化

Figure 8: 随机初始化 vs 零初始化的注意力层

关键发现: 随机初始化时，注意力 softmax 接近均匀分布，本质上执行因果均值池化，导致 token 表示坍缩和 MoE 层路由高度不平衡。

解决方案: 将注意力输出初始化为零（RMSNorm gains 设为零），使模型初始行为类似于逐 token 的 FFN 堆叠，注意力层的跨 token 交互在训练过程中逐渐启动。

四、预训练

数据

数据管道

Figure 4: HTML 预训练数据处理管道

  • 数据量: 30T tokens 主预训练 + 3.55T tokens 中期训练
  • 数据来源: 公开和授权的人类生成数据（web、GitHub、书籍、学术论文、新闻、多语言文本）
  • 关键原则: 无合成数据，无第三方蒸馏数据，努力避免和移除 AI 生成内容
  • 无开源训练数据集: 常见 ML 数据库从训练数据中去污

训练阶段

阶段Tokens上下文长度GB200 GPUs
预训练30T16,3848,192
中期训练 13.4T65,5368,192
中期训练 2150B262,1444,096

训练超参数

参数值
优化器AdamW (β₁=0.95, β₂=0.925)
权重衰减0.1（注意力 0.01，嵌入 0.005）
梯度裁剪1.0
学习率预热 12B tokens，余弦衰减 2e-4 → 2e-5
Dropout0.15
批大小134M tokens
数值精度BF16 权重/激活，FP8 GEMM，FP32 敏感位置

预训练损失

预训练损失

Figure 9: MAI-Base-1 在 30T tokens 上的预训练损失曲线

损失从 ~2.5 降至 ~0.87，多次早期阶段尖峰后快速恢复，无需人工干预。

基础模型对比

BPB 对比

Figure 10: 基础预训练模型的 Bits-per-byte 对比

MAI-Base-1 在所有四个 held-out 任务上以相似活跃参数数取得最低 BPB：

  • 优于 Gemma4 (31B)、DeepSeek V3.2 (37B/685B)、Kimi-K2 (32B/1T)
  • 仅 DeepSeek V4 Pro (49B/1.6T) 以 1.4x 活跃参数超越

分布式训练框架 YOLO

YOLO (You Only Launch Once) 是微软的内部大规模训练框架：

  • 内核: Triton、CUDA、CuteDSL、CUTLASS 自定义内核
  • 并行: 数据、张量、上下文、专家、流水线并行
  • 确定性: 位级可复现性
  • 容错: 分布式检查点 + 快速故障转移到热备
  • MFU: 持续保持 >20% MFU

五、强化学习

RL 总览

RL 总览

Figure 12: RL 训练总览

三个领域专家模型通过 RL 训练：

  1. STEM 推理专家: 数学、科学、竞赛编程
  2. Agentic 编码专家: SWE-Bench、工具使用
  3. 有用性和安全性专家: 人类偏好、安全信号

通过 Trace Distillation SFT 合并为统一模型，再进行最终 RL 阶段。

RL 目标

基于 GRPO 的 token-level policy gradient：

J(θ)=Eq,y1:G∼πold[1∑i=1G∣yi∣∑i=1G∑t=1∣yi∣min⁡(ri,t(θ)Ai,clip(ri,t(θ),1−ϵ,1+ϵ)Ai)]\mathcal{J}(\theta) = \mathbb{E}_{q, y_{1:G} \sim \pi_{\text{old}}} \left[\frac{1}{\sum_{i=1}^{G} |y_i|} \sum_{i=1}^{G} \sum_{t=1}^{|y_i|} \min\left(r_{i,t}(\theta) A_i, \text{clip}(r_{i,t}(\theta), 1-\epsilon, 1+\epsilon) A_i\right)\right]

自适应熵控制: 动态调整上界裁剪以维持目标策略熵

k←clip(k+δ⋅sign(H∗−H^(πθ)),0,kmax⁡)k \leftarrow \text{clip}\left(k + \delta \cdot \text{sign}(H^* - \hat{H}(\pi_\theta)), 0, k_{\max}\right)

奖励设计

R(q,yi)=Rtask(q,yi)+wlang⋅Rlang(yi)−wlen⋅Rlen(yi)R(q, y_i) = R_{\text{task}}(q, y_i) + w_{\text{lang}} \cdot R_{\text{lang}}(y_i) - w_{\text{len}} \cdot R_{\text{len}}(y_i)

  • 任务奖励: 基于代码执行或 AI 评审
  • 语言一致性奖励: 抑制 CoT 中的外语 token
  • 长度惩罚: 基于问题难度的自适应惩罚

采样策略

  • 问题采样: 早退策略 + 二次通过率过滤
  • Rollout 采样: top-p 采样，排除 nucleus 外 token
  • 长度课程: 从 8K 逐步增加到 128K tokens

自蒸馏

从 RL rollout 收集推理轨迹，对中期训练检查点执行 SFT：

  • O(1M) 推理轨迹足以匹配教师性能
  • 用于恢复运行失败、迁移到新基础模型、过滤 reward hacking
  • 轨迹应来自多个强检查点以增加多样性

STEM 训练

STEM 训练

Figure 15: AIME 2025 和 LiveCodeBench v6 上的 STEM 训练进展

RL 性能

Figure 1: MAI-Thinking-1 的 RL 训练性能

Agentic 训练

  • 多步骤编排的 agentic 循环
  • SWE 环境使用 bash 和 string-replace 工具
  • 前缀缓存命中率达 97-98%

RL 基础设施: Rocket

Rocket 框架

Figure 19: Rocket 大规模 RL 框架

  • 推理与学习器比例可达 5:1
  • 最大 RL 作业: 4864 GB300 芯片（4096 推理 + 768 学习器）
  • 基于 SGLang 的推理栈

六、评估结果

STEM 和 Agentic 编码

基准MAI-Thinking-1Sonnet 4.6Opus 4.6DeepSeek V3.2DeepSeek V4
AIME 202597.095.699.893.1—
AIME 202694.5————
HMMT Feb 202684.9———95.2
GPQA Diamond84.289.991.382.490.1
LCB v687.7——83.393.5
SWE-Bench Verified73.579.680.873.180.6
SWE-Bench Pro52.8—53.4—55.4
Terminal-Bench 2.046.059.165.446.467.9

通用能力

类别基准MAI-Thinking-1Sonnet 4.6
知识MMLU Pro8587
知识SimpleQA Verified3129
指令跟随IF Bench6950
指令跟随Adv. IF8586
长上下文GraphWalks (≤128k)9096
工具调用BFCL v37276
安全AIR-Bench8888
诚实TruthfulQA8888
健康HealthBench Prof.3538

人类评估

维度vs Sonnet 4.6vs Opus 4.6
整体偏好+0.07 ± 0.06-0.07 ± 0.06
简洁性和相关性+0.11 ± 0.02+0.07 ± 0.02
风格和语调+0.08 ± 0.02+0.05 ± 0.02
指令跟随-0.01 ± 0.02-0.04 ± 0.02
事实性-0.02 ± 0.02-0.03 ± 0.02

关键发现: 人类评估者偏好 MAI-Thinking-1 而非 Sonnet 4.6（49% 胜 vs 45% 负），但在简洁性和风格上明显优于两者。

安全评估

安全-有用性权衡

Figure 20: 安全-有用性权衡

越狱攻击成功率

Figure 21: 越狱攻击成功率

攻击类型MAI-Thinking-1Sonnet 4.6Opus 4.6
基础技术4.4%5.7%3.0%
组合技术17.6%15.0%17.4%
自适应技术26.8%26.4%25.1%

七、核心创新

创新点说明理论/实验依据
从零构建无蒸馏完全使用人类数据训练，无第三方模型蒸馏保证能力的可引导性和鲁棒性
高稀疏 MoE + Dense 交替8/512 experts + dense FFN 交替比均衡稀疏度更高效
LatentMoE压缩潜空间路由减少 all-to-all 通信开销
注意力零初始化RMSNorm gains 设为零消除初始 MoE 路由不平衡
自适应熵控制动态调整 RL 裁剪上界避免熵爆炸或坍缩
Outer Ratio Clip硬外裁剪防止梯度爆炸更稳定的 RL 训练
自蒸馏RL rollout → SFT → 继续 RL恢复失败、迁移能力
YOLO 框架确定性、容错的分布式训练90% goodput @ 8K GPUs
Dropless MoE完全无 token dropping避免因果泄漏和容量限制
全局 batch 负载均衡跨 DP workers 和 micro-batches 聚合比负载均衡损失类型更重要

八、训练基础设施

集群环境

  • 硬件: 8K GB200 GPUs，单一逻辑集群
  • 训练框架: YOLO (You Only Launch Once)
  • MFU: >20% 持续保持
  • Goodput: 90.0% @ 8K GPUs
  • 总开销: 51 小时（其中 6.5 小时重计算，14 小时非步进时间，18 小时 MFU 下降）

数值精度策略

位置精度
模型权重和激活BF16
前向 GEMMFP8 E4M3
数据梯度FP8 E5M2
权重梯度BF16
梯度累积FP32
残差流FP32
嵌入、RMSNorm、Router 权重FP32
优化器状态FP32

推理部署

  • 在 Microsoft MAIA-200 硬件上部署
  • 比 GB200 部署提高 40%+ token 生成吞吐量（相同机架功率预算）

九、总结

核心贡献

  1. 从零构建的 35B/1T MoE 推理模型: 无蒸馏，完全使用人类数据
  2. AIME 2025 达 97.0%: 超越 Sonnet 4.6 (95.6%)
  3. SWE-Bench Pro 达 52.8%: 接近 Opus 4.6 (53.4%)
  4. “爬山机器” 方法论: 将模型开发视为系统级优化问题
  5. YOLO 分布式训练框架: 确定性、容错、高 goodput
  6. 自适应 RL 熵控制: 动态维持策略熵的简单有效方法
  7. 安全-有用性平衡: 在多个安全类别上超越 Sonnet 4.6

技术影响

  • 无蒸馏范式: 证明从零训练可以达到与蒸馏模型相当的水平
  • MoE 架构: 高稀疏度 + dense 交替的有效性
  • RL 基础设施: 大规模 RL 训练的工程实践
  • 安全训练: 安全-有用性权衡的系统方法

局限性

  • Terminal-Bench 2.0 性能较低（46.0% vs Sonnet 59.1%）
  • 未使用 Terminal 交互环境训练
  • 多语言安全鲁棒性仍需改进
  • 108 页技术报告的复杂性

十、参考资源