MiMo-V2-Flash Technical Report

1 · 动机:为什么这一代 MoE 都奔向 hybrid attention

1.1 历史脉络:全 softmax 在长上下文时代的"两个崩塌"

2017 vanilla Transformer 之后,前 5 年 attention 的"全连接 softmax"几乎被默认是不可触碰的圣物。但到了 reasoning + agentic 时代(2024 起),两件事让 full softmax 不可持续:

• 序列长度爆炸。从 chat 时代的 2k–8k 到 reasoning 时代的 32k–256k,有些 agentic loop 一轮要喂 100k+ tokens。Attention 的 O(L²) 计算 + 线性增长的 KV cache,让 decode 时带宽(读 KV)而非计算成为新瓶颈。
• MoE 让模型变扁。MoE 把 FFN 拆成 256 个 expert 只激活 8 个,模型计算量被削到 1/N,attention 计算量在相对占比上反而越来越大。在 MiMo-V2-Flash 这种 309B/15B 设定下,decode 时 attention 占 step 时间的大头。

所以 2024–2026 这一波 MoE 几乎集体押注"sub-quadratic attention":

• DeepSeek-V3.2 走 sparse attention (NSA / DSA 类) — 改 score 函数
• MiniMax-01 / M1 走 lightning attention (linear attention 7:1 mix) — 改算子
• Gemma 2/3, gpt-oss, Qwen3-Next, MiMo-V2-Flash 走 SWA + Global hybrid — 保留 softmax 不变,只是把大多数层的 attention 范围切到 W token 内

Hybrid SWA 这条路最"保守"——softmax 计算公式不动,kernel 不用重写,只是 mask 形状变了——但工程友好度最高。MiMo-V2-Flash 押的就是这条。

1.2 别的方案为什么"够呛"

1.3 但"hybrid SWA 不平凡"在哪

看 Gemma 3 报告会发现一个反复出现的警告:SWA 的 window size 和 SWA:GA ratio 一旦激进,长 context 性能会严重塌方。Gemma 3 用 W=1024, ratio 5:1 才敢声明"性能不掉"。

MiMo-V2-Flash 的胆子:

• 把 W 从 1024 压到 128 — 缩了 8 倍
• SWA:GA 仍然保持 5:1
• 但能在 256K context 上 NIAH-Multi 拿 96.7%

能这么做的关键是引入 attention sink bias(从 gpt-oss 学来),它给每个 head 一个可学习的"沉没"通道,允许 token 在没有合适 key 时把注意力倒进 sink 而不是被迫均摊到所有 key。论文 ablation(Table 2)显示去掉 sink bias 这套架构是塌的——这正是论文最关键的"non-trivial"点。

同时,W=128 给出的是一种"分工压力":SWA 层没有能力处理远距离信息,被迫只做 local;远距离 dependency 必须由 9 个 GA 层(占总层数 ~19%)统一吸收。论文称之为 "clearer division of labor"——这是个有点诗意但工程上确实成立的解释:模型不再有"两边都不专"的层。

2 · 背景速查

2.1 关键术语

2.2 标准回忆: SWA 的 mask 长什么样

3 · 架构 1: 5:1 Hybrid SWA + Global Attention

3.1 架构总览

整个 backbone 共 48 层:

• 第 1 层(特殊):Global Attention + dense FFN(不是 MoE),起到 "stabilize early representation" 的作用
• 第 2–48 层:组成 8 个 Hybrid Block。每个 block = 5 个 SWA 层 + 1 个 GA 层,所有 FFN 都是 sparse MoE

合计: 39 SWA + 9 GA = 48 层 (其中 GA = 1 dense + 8 在每 block 末尾)。

3.2 数字感受 — KV cache 真的省了多少?

设序列长 L=256K, head_dim_v=128, num_kv_heads=8 (取 SWA 设置), 每 KV slot in fp8 ≈ 1 byte/element。

缩减 ~5.3×(论文宣称 ~6×, 略乐观)。SWA 层贡献几乎可忽略 — KV 主要由 9 个 GA 层占。

3.3 反向论证: 为什么 W=128 比 W=512 反而好?

论文 Table 3 实测:

• W=128 + sink: GSM-Infinite 17.3, NoLiMa 51.2, MRCR 34.4
• W=512 + sink: GSM-Infinite 17.2, NoLiMa 38.5, MRCR 19.6

反直觉。论文给出的解释:"clearer division of labor"。W=128 时 SWA 层物理上无法处理远距离依赖,模型被迫把"远距离任务"全部交给 GA 层。W=512 时 SWA 层偶尔也能看到 256–512 距离的 token,变得"半懂不懂",最终两类层都没专精。

这听起来像后验解释,但和稀疏专家网络中"specialization"的现象学是一致的——给一个 expert 能力上限,反而能逼出更清晰的角色。

4 · Attention Sink Bias: 让 W=128 不掉点的关键

4.1 公式

原版 attention:

加 sink bias 之后:

每个 head 一个标量 sink ∈ ℝ(可学习)。注意 sink 项加在分母里,但不对应任何 value——也就是说,如果 sink 很大,所有 s_ij 一起被压小,等效于"这个 token 的 attention 总和 < 1",可以"什么都不看"。

4.2 直觉:为什么 SWA 离不开 sink

SWA 的天然缺陷:窗口 W=128 内,如果当前 token 的 query没有任何真正相关的 key,标准 softmax 会强制把 attention 均摊到这 128 个 key 上。模型没有"我什么都不需要"的选项。

这在 long-context 场景非常致命: token 间距离一旦 > W,模型完全应该"忽略"窗口内的旧 key,但 softmax 不允许。

Sink bias 提供了一个"逃生通道":当 q_i · k_j 都很弱时,sink 项主导分母,所有 s_ij ≈ 0,output ≈ 0(value 加权和接近 0 向量)。模型可以学习"把消息留给残差/下一层处理"。

4.3 Worked example

设某 head 在第 5000 个 token 上,query q 在 W=128 窗口内的 128 个 key 上得到 logits 全部都在 0~2 之间(没有强信号),sink 学到 = 4。

• m_i = max(max(0~2), 4) = 4
• 分子项: exp(a_ij − 4) ≤ exp(-2) ≈ 0.135 (每项)
• sink 项: exp(4 − 4) = 1
• 分母 ≈ 1 + 128 × 0.135 ≈ 18.3
• 每个 s_ij ≤ 0.135 / 18.3 ≈ 0.0074;128 个加起来 ≈ 0.945,但 sink "吃" 0.055

看起来吃得不多,但实测里 sink 学到的值随 head 不同从 0 到 6 都有。sink 大的 head 显示出"有事就用,没事关掉"的开关性质——这是 attention sink 论文(Xiao et al. 2023)早期观察到的现象的"显式参数化"。

5 · 架构 2: 256 专家细粒度 MoE

5.1 设计选择

MiMo-V2-Flash 没有 shared expert —— 这跟 DeepSeek 系列正好相反。Shared expert 的本意是"all-token 共有的基础能力",但实测中容易让模型把通用知识塞给 shared,稀释 routed expert 的特化。MiMo 选择不要这层"安全垫",彻底 rely on 256 个 routed expert。

5.2 Auxiliary loss & expert bias

MiMo-V2-Flash 用两层平衡机制:

1. Sequence auxiliary loss(系数 1e-5):标准 load-balance loss
2. Expert bias update(0.001 in stage 1/2, 1e-5 in stage 3):每个 expert 一个动态 bias 加到 routing score 上,根据负载实时调整。和 DeepSeek-V3 的 auxiliary-loss-free balancing 是同一个 trick。

5.3 SFT 期间的"num-zeros 监控"

论文有个有意思的工程发现 — SFT 阶段最稳定的训练健康指标不是 loss,而是 zero-gradient 参数比例:

• num-zeros 上升 → expert 负载失衡(部分 expert 被路由 0 token,梯度全为 0)
• num-zeros 下降 → 模型在过拟合(梯度 dense 到处都在更新,通常意味着 catastrophic memorization)

调它的两个 lever:expert bias update rate 和 AdamW ε。这种"看似 trivial 的 SFT 工程"在 MoE 上经常决定后续 RL 能不能跑起来。

6 · 架构 3: 轻量 MTP — 同时是训练 aux 和 draft

6.1 MTP 是什么

Multi-Token Prediction (Gloeckle et al. 2024):标准 LM 训练只预测下 1 个 token,MTP 让模型同时预测下 1, 2, ..., K 个 token。每个未来位置一个独立的 head,共享主干 hidden states。

6.2 本文的设计取舍

6.3 关键工程取舍

6.4 worked example: 一次 spec decoding step

设当前已生成 token 序列 "3 + 4 = ",token id 序列 [..., 91]:

1. 主模型 forward 一步,得 last hidden h_t ∈ ℝ⁴⁰⁹⁶,主模型自身预测下 1 token = "7" (作为已采样的 t+1)
2. MTP head 1 输入 [h_t, emb("7")],SWA + dense FFN → predict t+2 = " " (空格,概率 0.95)
3. MTP head 2 接 head 1 的 hidden + emb(" ") → predict t+3 = "答案" (概率 0.31, 不算很高)
4. MTP head 3 接 head 2 的 hidden + emb("答案") → predict t+4 = "是" (概率 0.40)
5. 主模型在下一次 forward 时,把 [t+1, t+2, t+3, t+4] 一次输入,得到每个位置的真实 logits
6. 对 t+2/t+3/t+4 逐个 rejection sample:接受 t+2 概率 ≈ 0.95(几乎肯定),接受 t+3 概率 = min(1, p_main("答案")/0.31) — 假设 p_main = 0.10,接受概率 0.32,采样到拒绝就停
7. 本步合计产出 1 (主) + 1 (t+2 接受) + 0 (t+3 拒绝, 重采) = 3 个 token,acceptance length α = 3

注意 MTP 的 spec decoding 与 EAGLE-3 的本质区别:MTP 是"模型自带的",不需要单独训练一个外部 draft 模型;EAGLE-3 是 add-on draft head,要做单独的 distillation。MTP 等于把 spec decoding 内化进 pre-training。

6.5 实测数字 (论文 Figure 7 + Table 10)

• WebDev (低 entropy): α = 3.6 — 3 个 head 都接受意味着每步 4 token
• MMLU-Pro (高 entropy): α ≈ 2.9
• 实测 decoding speedup: bs=64, α=3.8, 2.67×;bs=96, α=3.8, 2.70×

Acceptance length 与 next-token entropy 的拟合: α = 4(1 − 0.58·entropy^0.58),R²=0.995 — 这个拟合相当干净,说明 MTP 的 acceptance 几乎完全由模型自己的 prediction 信心决定,与具体任务的"难度"无关,只与 entropy 相关。

7 · 预训练: 27T tokens 三阶段 + FP8

7.1 三阶段 schedule

7.2 Learning rate schedule

• Stage 1: linear warmup 0 → 3.2e-4 (over 50B tokens) → constant 3.2e-4 for 12T → cosine decay to 1.0e-4 over 10T
• Stage 2: 1.0e-4 → cosine decay to 3.0e-5 over 4T
• Stage 3: 3.0e-5 → cosine decay to 1.0e-5
• Batch size: warmup to 2048 over first 500B, hold; Stage 3 用 256

7.3 MTP loss weight 调度

MTP aux loss 占比有意识地 anneal:Stage 1 = 0.3, Stage 2/3 = 0.1。直觉是早期需要靠 MTP 的"前瞻信号"提升表征质量,后期主任务已经够强,MTP 转为"配角"。

7.4 Pre-train 数据细节(论文披露不多)

训练语料:web pages + 书 + 学术论文 + code + math + STEM。重点:

• "deliberate shift toward data exhibiting long-range dependencies" — 长篇 web 文档 + 仓库级代码 + PR/issue/commit 历史
• 合成 reasoning data ~5% 在 mid-train 引入(没说细节,大概是 CoT 风格的数据合成)
• 没有公开数据混合比例,也没有 quality filter 细节

7.5 FP8 mixed-precision

主体 FP8(类似 DeepSeek-V3),但保留:

• BF16: attention output projection, embedding, output head
• FP32: MoE router parameter

这套 FP8/BF16/FP32 的精度切割是当前 MoE 训练的"通用方案",DeepSeek-V3 之后基本是 production 标配。

8 · 后训练: MOPD (Multi-Teacher On-Policy Distillation)

这是论文最有"研究气"的部分。

8.1 问题:capability imbalance + learning inefficiency

标准做法是"先 SFT 然后 RLHF",或者跑多个独立 RL 然后参数合并。两条路都有问题:

• See-saw 效应:提一项掉一项(数学好了写作差了)
• 合并参数损能力:weight averaging / TIES 这类 model merging 在 frontier 模型上保留率不高
• 静态 distillation 数据 → 分布漂移:teacher 的 trajectory 不是 student 自己生成的,exposure bias 严重

8.2 MOPD 三阶段

8.3 Loss 推导

从 reverse KL 出发(Eq. 5):

对 θ 求梯度(Eq. 6,REINFORCE-like):

把 log(π_T/π_θ) 当作"per-token advantage"。再加上 train/inference distribution mismatch correction(因为 inference 用 μ_θ 采样,training optimize π_θ),得到 surrogate(Eq. 7-8):

三件事融合在一个 loss 里:

1. w_t: importance sampling 的 truncation,过滤 train/inference mismatch 异常大的 token (gpt-oss / Zhao 2025 套路)
2. log(π_T/π_θ): 蒸馏信号,token-level dense
3. α · Â_ORM: 经典 outcome reward (GRPO 算的 advantage)

8.4 实测效果 (论文 Table 7)

大部分项 student 能匹敌甚至略超 best teacher;但 BrowseComp 和 creative writing 出现明显 regression — 论文承认 MOPD 不是 silver bullet。我个人解读:MOPD 适合"有可验证 reward 的 hard skill 域",对开放式生成 / 长 horizon search 仍然不稳。

9 · RL 基础设施: R3 / Data Scheduler / Toolbox

9.1 Rollout Routing Replay (R3) — 解决 MoE rollout/train mismatch

MoE 训练的一个老问题:rollout 用 SGLang(inference engine, FP8 计算可能不严格 deterministic),training 用 Megatron-LM,两边对同一个 token 算出的 router score 不完全一致 → expert 选择不一致 → π_θ(y_t) 在 training/inference 不同 → importance sampling 假设崩塌 → RL 训不稳。

R3 的解法很直接:training 时强制使用 rollout 时记录下的同一份 expert routing 决策。这意味着 routing 不再是 training 的一部分 — 它被"freeze"在 rollout 决策上。开销的关键在压 routing log 的存储/通信成本。

9.2 Data Scheduler — 解决 long-tail straggler

RL rollout 经常出现"绝大多数轨迹 1k tokens 内结束,少数 32k 不收"的长尾分布。同步等会让 GPU 空转。本文方案:

• partial rollout:把超长轨迹分多 step 处理,每 step 限制 partial sample 比例避免 staleness 过大
• staleness-aware truncated importance sampling:partial 轨迹横跨多个 policy 版本,用 truncated IS 修正
• priority-based scheduling:按 data source 配比 + 历史 pass rate 动态分配新 prompt 到空闲 GPU

9.3 Toolbox — 解决 agentic RL 的 tool 资源争用

• centralized resource allocator,enforce QPS 限制
• 用 Ray actor pool 避免 cold start
• environment pre-warming + sequence-level async reward

这一节细节不多,但传达的信号是 agentic RL 的瓶颈早就不是 GPU 算力,而是外部 tool 的资源调度。

10 · MTP 推理加速实测

10.1 Acceptance ↔ Entropy 拟合

论文 Figure 7 的拟合公式:

其中 H = next-token cross-entropy(模型自己的预测信心)。

10.2 Decoding speedup (论文 Table 10)

16K input, 1K output。bs=96 是甜点 — 太小没填满 SM,太大 attention KV I/O 成瓶颈。

11 · 与同类工作对比

11.1 几条关键差异化

• vs DeepSeek-V3.2: DSv3.2 用 Dense Sparse Attention (DSA) 替代 GA, 长 context retrieval 性能差(Table 6 长 context 严重塌方);MiMo 走 SWA 保留全 softmax 语义,长 context 反而更鲁棒
• vs Qwen3-MoE: 都是 256 routed (近似), 但 MiMo 没有 shared expert,且 SWA 比 Qwen3 激进得多
• vs MiniMax-M1: linear attention 路线在 long retrieval 仍偏弱,MiMo hybrid 是更安全的路
• vs Gemma 3: Gemma W=1024,MiMo 直接押 W=128 + sink — 同样 5:1 的 ratio 下,KV cache 8× 缩减

12 · 局限 / 个人 take / 待验证问题

论文承认的局限

• 架构探索"preliminary",没系统 sweep 不同 ratio / window size 的组合
• 仍逊于最强 closed model (Gemini 3 Pro, GPT-5 High) — knowledge capacity 是 SimpleQA 上 20.6 vs Kimi-K2-Base 35.3 的明显短板
• MOPD 的 teacher 训练成本不算在内,实际管线计算量比传统 RL 大

我的疑问

1. W=128 极限究竟在哪?论文只 sweep 了 128 vs 512。W=64 / 32 加 sink 可能依然能 work,这意味着 KV 还能再缩。
2. Sink bias 学到什么?论文没分析每个 head 学到的 sink 值分布。GA 层的 sink 是不是接近 0(因为 GA 不需要 escape valve)?
3. MOPD 在 BrowseComp / creative writing 上掉点 — 是不是 token-level reverse KL 在长序列复杂任务上 over-regularize 了?把 KL 换成 forward KL 或 JS 会不会更好?
4. R3 是不是过度保守?training 时强制用 rollout 的 expert routing,等于让 router 不再被 RL gradient 影响。这是不是变相冻结了 routing 的演化?
5. 3-layer MTP 的内存代价:K=3 时每个 head 0.33B,推理时还要维护 3 套独立的 SWA KV cache。在长 context 下,这块 KV 会不会再次成为瓶颈?
6. 27T tokens 数据细节非常 sparse — 没说预训练数据混合比例,没说 quality filter。复现性接近 0。
7. SWE-Bench Verified 73.4% 实测有 reward hacking(Appendix B 提到),论文没披露去除 hacking 后的真实分数。

记忆点

精读笔记 · 配套论文 PDF: /data/szhang967/papers/paper-notes/MiMo-V2-Flash_2601.02780.pdf