跳转至

Mode-conditioning unlocks superior test-time compute scaling

会议: ICLR 2026
OpenReview: https://openreview.net/forum?id=JzkdJQzPw1
代码: 待确认
领域: LLM 推理 / 测试时计算扩展
关键词: 并行采样, 多样性坍缩, Pass@k, 模式条件化, 测试时扩展, 蒸馏

一句话总结

针对并行采样中"模型坍缩到单一推理策略、重复采样反复犯同样错误"的多样性坍缩问题,本文提出模式条件化(ModC)框架——用专家模型或模式前缀显式地把测试时算力均匀分配到不同推理模式上,从而在数学推理与图搜索任务上把 Pass@k 扩展曲线整体抬高,并带来约 4× 推理效率提升。

研究背景与动机

  • 领域现状:并行采样(给模型同一题独立采 k 次取最优)是测试时扩展和 RL 的支柱,在数学、代码、科学发现等可自动验证的领域尤其有效,其性能由 \(\text{Pass@}k_{\text{std}}(x)=1-(1-p_x)^k\) 刻画。
  • 现有痛点:SFT 和 RL 都会引发多样性坍缩(diversity collapse)——模型集中到少数几个模式上,导致多采的样本只是把同一个错误再犯一遍,算力一加上去就边际递减。在某些题上成功策略的概率 \(p_x\) 被压得极小,需要不切实际的采样量才能解出。
  • 核心矛盾:以往缓解坍缩的做法(Pass@k 训练、权重正则)大多仍依赖逐 token 的温度采样来制造多样性,而温度采样改变的是局部 token 分布,无法在"高层推理模式"这个尺度上保证覆盖。作者用 Countdown 任务说明:一道题可能 DFS 能解 BFS 不能(反之亦然),但标准模型往往对单题只稳定采到一种模式,押错就全军覆没。
  • 本文目标:不去修改微调目标来"防坍缩",而是走互补路线——显式地把采样预算分摊到多个模式上,强制覆盖既包含主导策略也包含可能在主导策略失败处成功的备选策略。
  • 核心 idea【模式条件化】 把测试时扩展围绕"多个推理模式"来组织,并给出可在实践中可靠采到指定模式的两种训练方式(专家模型 / 模式前缀);当数据没有现成模式标签时,用梯度聚类自动发现模式。

方法详解

整体框架

ModC 的出发点是一条简单不等式:与其从模型自己那个"不确定混合"分布里反复采样,不如把 k 次预算均匀切给若干已知模式。训练阶段把数据按模式划分并让模型学会"按指令切模式"(专家模型或前缀两种实现),推理阶段对每个模式均分采样预算。当模式标签不存在时,先用梯度聚类把训练样本分成 C 个"模式",再套用同一套 ModC 流程。

flowchart TD
    A[训练数据] --> B{是否有模式标签?}
    B -- 有(搜索算法/教师身份) --> C[按模式划分数据]
    B -- 无 --> D[梯度聚类<br/>∇θ log pθ y x → 随机投影 → C 簇]
    D --> C
    C --> E1[实现一: 训练独立专家模型]
    C --> E2[实现二: 单模型 + 模式前缀 Mode k]
    E1 --> F[推理: k 次预算均分到各模式]
    E2 --> F
    F --> G[Pass@k 整体抬升 / ~4× 效率]

关键设计

1. 均分模式预算严格优于自由混合采样:理论保证。 ModC 的合理性首先来自一个可证明的优势。设两个模式在输入 \(x\) 上的成功概率为 \(p_{1,x}, p_{2,x}\),均分预算后整体成功率为 \(\text{Pass@}k_{\text{ModC}}(x)=1-(1-p_{1,x})^{k/2}(1-p_{2,x})^{k/2}\)。反过来,若模型不知道哪个模式更好,以随机权重 \(w_x\)(满足 \(\mathbb{E}[w_x]=1/2\))混合两模式,则单样本成功率是 \(w_x p_{1,x}+(1-w_x)p_{2,x}\)。由于该函数对 \(w\) 是凹的,Jensen 不等式给出 \(\mathbb{E}_{w_x}[\text{Pass@}k_{\text{std}}(x;w_x)]\le \text{Pass@}k_{\text{std}}(x;1/2)\);而只要 \(p_{1,x}\ne p_{2,x}\),就有 \((1-p_{1,x})^{k/2}(1-p_{2,x})^{k/2}<(1-p_x)^k\),于是 \(\text{Pass@}k_{\text{ModC}}(x)>\text{Pass@}k_{\text{std}}(x;1/2)\)。结论是:哪怕模型的模式偏好是任何以 0.5 为中心的分布,显式均分都严格更优——这正是 ModC 不靠改目标、只靠"分配算力"就能赚到 Pass@k 的根因。

2. 专家模型实现:强特化、低相关错误。 最直接的落地是把训练数据按策略切成子集,在每个子集上各训一个独立模型(总数据量与计算量保持不变),测试时把预算均分给各专家(两模式即各采 \(k/2\))。这种切分天然让每个专家高度特化、彼此犯的错误更不相关,从而把并行采样的收益放大。代价是模式之间无法共享知识——在 Countdown 这种模式本身差异大、共享需求低的任务上,专家模型反而比前缀更强(Pass@1024 在对抗集上最高拉开约 20% 的差距)。

3. 模式前缀实现:单模型内特化、跨模式共享知识。 为弥补专家模型不能共享知识的缺陷,作者借鉴可控文本生成里的条件 token 思路,在输入前拼上离散条件标记(如 [Mode 1][Mode 2] 或教师身份 token),训练模型把每个前缀绑定到一种推理策略;推理时通过对前缀均匀采样来强制均衡分配。这样模型既能特化出不同模式,又能在模式间共享语言/数学基础——在数学推理(短/长 CoT)这类策略共享需求高的任务上,前缀普遍优于专家模型,印证了"知识共享在数学任务里更关键"。

4. 梯度聚类自动发现模式:摆脱标签依赖。 现实数据很少自带 DFS/BFS 或教师身份这类清晰标签。作者对每个训练样本 \((x,y)\) 计算梯度 \(g_\theta(x,y)=\nabla_\theta \log p_\theta(y|x)\),用 Rademacher 随机投影降维后做 K-means 聚类成 C 簇,把同簇样本视作同一"模式"再套用 ModC。在 NuminaMath 上这种无任何额外信息的自动发现也能稳定涨 Pass@k(最高约 10%),说明标准训练确实没把数据里隐含的多样性用满。

实验关键数据

主实验表格

设置 任务/评测 基线对比 ModC 收益
Countdown 图搜索(自然集) Pass@1024 标准训练 专家模型最高 +8%
Countdown 图搜索(对抗集,单算法可解) Pass@1024 标准训练 最高 +20%
短 CoT 多教师蒸馏(NuminaMath→MATH500) Pass@k 混合教师/最佳单教师 Qwen2.5-0.5B +10%,OLMo2-7B +15%(前缀更优)
长 CoT 多教师蒸馏(OpenThoughts→AIME2025) Pass@k QwQ-32B/DeepSeek-R1 单教师 超过最佳单教师;用 k=256 匹配标准训练 k=1024,约 4× 效率
自动模式发现(NuminaMath→MATH500) Pass@k 标准训练 跨 0.5B–7B 一致提升,最高约 +10%

消融实验表格

消融维度 观察结论
训练数据平衡性(拒绝采样偏向 DFS vs 50-50) 即使用平衡数据,标准训练仍对很多题做出极不均衡的模式分配;ModC 才把每题 BFS 比例集中到约 0.5
模型规模(0.5B → 7B,Qwen2.5 / OLMo2) ModC 增益在所有规模一致出现
专家模型 vs 前缀 Countdown(模式差异大)专家更优;数学(需共享知识)前缀更优
随机划分对照 把数据随机分两组有时也有收益,但不及按真实模式的 ModC
ModC + RL / Pass@k RL RL 把两者 Pass@1 拉到同一水平,但 ModC 在 k=2 即领先;在显式防坍缩的 Pass@k RL 之上仍能再加分

关键发现

  • 多样的训练数据只有配上"保模式"的机制才有用:标准混合教师常常打不过最佳单教师(反直觉),而 ModC 能把教师多样性真正兑换成更强的测试时扩展。
  • ModC 富化解空间但不牺牲 top 输出:与标准 SFT 不同,RL 后 ModC 的 Pass@1 不掉,却在 k≥2 立刻领先。
  • 效率提升直接可量化:长 CoT 上 ModC 用 1/4 的样本数就匹配标准训练的 Pass@1024。

亮点与洞察

  • 把"多样性"从 token 尺度抬到模式尺度:温度采样调的是局部分布,ModC 直击高层策略覆盖,思路简洁却抓到了并行采样失效的真正瓶颈。
  • 理论与工程双落地:Jensen 不等式给出"均分严格更优"的干净证明,两种实现都是工程上轻量、易接入现有蒸馏/SFT 流程的改动。
  • 梯度聚类把方法从"需要标签"解放到"通用数据",大幅扩展了适用面,也反过来证明标准训练在浪费数据里的隐含多样性。
  • 与 RL 正交可叠加:ModC 是 SFT 侧干预,可叠在标准 RL 乃至专门防坍缩的 Pass@k RL 之上继续涨。

局限与展望

  • 前缀变体的 RL 仍是空白:作者担心 RL 会破坏"前缀↔模式"的绑定(需要前缀跟随奖励),故 RL 实验只做了专家模型变体,前缀+RL 留作未来工作。
  • 模式数与聚类粒度:梯度聚类的簇数 C、投影维度等超参对发现的"模式"质量影响较大,论文未给出系统化的选择准则。
  • 预定义模式仍带人工先验:Countdown 的 BFS/DFS、多教师身份都属于"恰好已知"的干净场景,更开放任务里"什么算一个模式"仍不清晰。
  • 均分预算的最优性:理论只证了"均分优于自由混合",并未论证均分就是各模式成功率未知时的最优分配,自适应分配可能进一步提升。

相关工作与启发

  • 改进并行测试时扩展:相比 Pass@k 训练、权重集成(Dang et al. 2025)、多样化 beam search、多样化提示等"防坍缩"路线,ModC 走的是数据中心的条件化路线,把专家模式显式编码进模型。
  • 专业化训练:与 Mixture-of-Experts 最接近,但 MoE 是把不同数据路由到模型的子部件以省激活参数,ModC 则是整模型处理所有数据、只在"输出模式"上做条件化,目标是覆盖多样策略而非省算力。
  • 创造力与多样性研究:呼应了"温度采样与创造力仅弱相关、还会引入不连贯"以及"全局规划/种子条件化对创造性生成关键"等发现,ModC 可看作在推理任务上对"种子/条件化提升多样性"的一次具体兑现。

评分

  • 新颖性: ⭐⭐⭐⭐ 把多样性问题从 token 尺度重构到"推理模式"尺度,并给出干净的理论保证,角度新且抓到要害;两种实现本身在别处出现过,故非满分。
  • 实验充分度: ⭐⭐⭐⭐ 覆盖图搜索+短/长 CoT、两大模型族、0.5B–7B、标签/无标签、SFT/RL 多设置,结论一致;但缺前缀+RL、最优预算分配等收尾实验。
  • 写作质量: ⭐⭐⭐⭐ 动机—理论—实现—实验线索清晰,图示直观;部分小节有笔误。
  • 价值: ⭐⭐⭐⭐ 改动轻、可叠加、效率收益可量化(~4×),对所有做测试时扩展/蒸馏的团队都有直接可用价值。