$S^3$: Stratified Scaling Search for Test-Time in Diffusion Language Models

Each language version is independently generated for its own context, not a direct translation.

这篇论文介绍了一种名为 S3（分层缩放搜索）的新方法，旨在让现有的“扩散语言模型”（一种新型 AI）在不重新训练的情况下，通过“多思考一会儿”来变得更聪明。

为了让你轻松理解，我们可以把 AI 写答案的过程想象成在迷雾中绘制一幅画，或者在迷宫中寻找出口。

1. 背景：AI 是如何“思考”的？

传统的 AI（比如现在的聊天机器人）像是一个按部就班的画家。它从左到右，一笔一划地写，写错一个字就很难回头改，只能硬着头皮继续画。

而这篇论文研究的扩散语言模型（DLM）则像是一个先画草稿再精修的画家。

过程：它一开始在一张全是乱码（迷雾）的画布上，然后一步步把乱码擦除，慢慢显现出清晰的文字。
问题：传统的做法是，AI 只画一条路径，从乱码到成品，一气呵成。如果第一步走错了，后面可能全错，而且它不知道自己在迷雾里是不是走偏了。

2. 痛点：为什么“多试几次”不够用？

以前，人们想让 AI 变强，常用的笨办法是Best-of-K（最佳 K 选）：

做法：让 AI 画 10 次（生成 10 条不同的路径），然后挑出最好的一张。
比喻：就像你在迷雾里派 10 个探险队，每个队都随机乱走，最后看谁运气好撞到了宝藏。
缺陷：这就像“撒网捕鱼”，虽然网撒大了，但所有探险队都是从同一个错误的起点（迷雾）出发的。如果迷雾本身的方向感就是错的，派再多的人也是徒劳。而且，随着人数增加，收益越来越小（边际效应递减）。

3. 解决方案：S3（分层缩放搜索）

S3 的核心思想是：不要等到最后才挑答案，而是在“画画”的每一步，都聪明地分配精力。

我们可以把 S3 想象成一支拥有“千里眼”的精英探险队：

第一步：分兵探路（扩展）

在迷雾中的每一步（比如擦除一点乱码时），S3 不会只让一个探险队走，而是让N 个小队同时出发。每个小队尝试走不同的方向（生成几个候选方案）。

第二步：千里眼评估（验证）

这是最关键的一步。S3 配备了一个轻量级的“验证员”（Verifier）。

比喻：验证员就像是一个拿着地图的向导。虽然向导看不到终点（因为还没画完），但他能根据当前的草图，快速判断：“嘿，这个方向看起来像是有宝藏的！”或者“那个方向全是死胡同，快回来！”
特点：这个向导不需要知道标准答案（Ground Truth），它只看逻辑通不通、数学对不对、结构像不像。

第三步：动态调整（重采样）

根据向导的反馈，S3 会把资源集中在好路上。

比喻：如果向导说“左边那条路有 80% 概率有宝藏”，S3 就会把右边那些走错路的探险队叫回来，让他们加入左边那条路，继续深入探索。
技术点：它使用了一种叫 SSP（Srinivasan 采样过程）的算法，确保在集中资源的同时，不会把所有鸡蛋都放在一个篮子里（保持多样性，防止大家都撞死在同一棵树上）。

第四步：最终决战

经过 T 步这样的“分兵 - 评估 - 集中”循环，原本分散的迷雾逐渐汇聚成几条清晰的大路。最后，S3 从剩下的几条好路中，选出最靠谱的一个答案。

4. 为什么 S3 很厉害？（核心优势）

不重新训练：就像给现有的汽车装了一个更聪明的导航仪，不需要重新造发动机（不需要重新训练模型）。
算得值：它把计算力用在了刀刃上（迷雾中的关键决策点），而不是盲目地多画几张图。
效果显著：
- 在数学题（MATH-500）上，准确率从 25.6% 提升到了 30.2%。这就像让一个原本只能做对 1/4 题的学生，通过这种“分步检查”的方法，做对了 3/10 题，进步巨大。
- 在常识问答和逻辑推理上也有明显提升。

5. 总结

如果把 AI 生成文本比作在迷雾中找路：

传统方法：蒙着眼走，或者派一群人蒙着眼乱撞，看谁运气好。
S3 方法：派出一支小队，每走一步就停下来，用“千里眼”看看哪条路更有希望，然后让所有人集中力量走那条路，同时保留一点探索其他可能性的机会。

这篇论文证明了，对于这种新型 AI，“边做边检查、动态调整方向” 比单纯的“多试几次”要高效得多。这为未来让 AI 更聪明、更可靠提供了一条实用的新路径。

Each language version is independently generated for its own context, not a direct translation.

1. 研究背景与核心问题 (Problem)

背景：
扩散语言模型（Diffusion Language Models, DLMs）通过迭代去噪过程生成文本，具有独特的结构优势。近年来，DLMs 在性能上已展现出与自回归模型（Autoregressive Models, AR）相当甚至更强的潜力。

核心问题：
现有的测试时扩展（Test-time Scaling）策略（如 Best-of-K 采样）在 DLMs 上存在根本性局限：

密度 - 质量不匹配 (Density-Quality Mismatch)： 基础扩散模型分布 $p_0(x)$ 的高概率区域往往与验证器（Verifier）评分高的高质量输出区域不一致。
Best-of-K 的局限性： 简单地增加采样数量 $K$ 并不能改变底层分布。根据极值理论，Best-of-K 带来的性能提升仅随 $K$ 对数增长（ $\log K$ ），效率低下。
未充分利用去噪结构： 与自回归模型不同，DLMs 的去噪过程是一个包含 $T$ 步的随机序列过程。标准解码通常只采样单条轨迹，忽略了在去噪过程中探索多条潜在轨迹以优化最终输出的机会。

目标：
在不重新训练模型、不改变解码调度（Decoding Schedule）的前提下，通过重新分配推理时的计算资源，使 DLMs 能够生成更高质量的输出。

2. 方法论：S3 (Stratified Scaling Search)

作者提出了 S3，一种基于验证器引导的粒子搜索方法。其核心思想是在去噪过程中动态地重新分配计算资源，将粒子（候选轨迹）引导至验证器评分更高的区域，从而近似一个“奖励倾斜（Reward-tilted）”的分布。

2.1 理论推导

最优目标分布： 在 KL 散度约束下，最大化期望验证器奖励的最优分布是吉布斯倾斜分布（Gibbs-tilted distribution）：
$\tilde{p}_0(x) \propto p_0(x) \exp(\tau f(x))$
其中 $f(x)$ 是验证器分数， $\tau$ 是温度参数。
路径空间问题： 将去噪过程视为路径空间问题。理想情况下，应通过 Doob h-变换扭曲反向转移核，但这需要计算不可行的后向信息函数（backward information function）。

2.2 S3 算法流程 (Level 3 近似)

S3 通过三个层级实现，最终采用一种可计算的近似方案：

初始化： 在 $t=T$ 时刻初始化 $N$ 个粒子（完全掩码序列）。
迭代去噪 (Expand-Score-Resample)： 对于每一个去噪步骤 $t$ $t$ （从 $T$ $T$ 到 $1$）：
- 扩展 (Expansion)： 将每个粒子 $x_t^{(i)}$ 扩展为 $b$ 个子候选者 $x_{t-1}^{(i,j)}$ ，通过基础模型采样得到。
- 前瞻评分 (Look-ahead Scoring)： 由于无法直接获得最终输出 $x_0$ ，S3 利用模型的单步去噪预测（One-step clean prediction） $\hat{x}_0^{(i,j,t)}$ 作为前瞻估计。使用轻量级、无需真值标签的复合验证器 $f(\cdot)$ 对预测结果进行评分 $s_{i,j,t}$ 。
- 重采样 (Resampling)： 根据评分计算重要性权重 $\tilde{w} \propto \exp(\lambda s)$ 。利用 Srinivasan 采样过程 (SSP) 进行低方差的依赖舍入重采样，将 $N$ 个粒子的预算分配给得分更高的候选轨迹。SSP 保留了随机性，避免了确定性 Top-K 剪枝导致的模式崩溃。
最终输出： 在 $t=0$ 时刻，对剩余的 $N$ 个粒子进行多数投票（Majority Voting），若平局则选择基础模型负对数似然（NLL）最低的序列。

2.3 验证器设计 (Composite Verifier)

验证器 $f(x)$ 无需 Ground Truth，基于生成文本的内在信号构建，包含五个正交维度：

结构完整性 (Structural Completeness)： 检查关键词、答案格式（如 \boxed{}）。
算术一致性 (Arithmetic Consistency)： 验证推理过程中的等式计算是否正确。
答案可达性 (Answer Reachability)： 检查最终答案是否在推理过程中被推导出来。
模型置信度 (Model Confidence)： 基于生成 token 的对数概率。
非退化性 (Non-degeneracy)： 惩罚重复、过短或信息量低的输出。
注：针对不同任务（如数学题 vs 选择题），各维度的权重 $\alpha_k$ 会动态调整。

3. 主要贡献 (Key Contributions)

识别密度 - 质量不匹配： 首次形式化地指出 DLMs 中基础分布 $p_0(x)$ 与验证器奖励分布之间的错位，并证明这限制了 Best-of-K 的效果。
理论目标： 证明了在 KL 约束下，最优推理目标应为奖励倾斜的吉布斯分布，并提出了通过去噪轨迹搜索来近似该分布的方法。
提出 S3 算法： 设计了一种无需重训练、基于验证器引导的粒子搜索方法。它利用轻量级验证器和 SSP 重采样，在去噪过程中动态调整计算资源。
实证效果： 在 LLaDA-8B-Instruct 模型上，S3 显著提升了多个基准测试的性能，特别是在数学推理任务上，且优于 Best-of-K 策略。

4. 实验结果 (Results)

实验基于 LLaDA-8B-Instruct 模型，在四个基准测试上进行评估：

数据集	基线 (Baseline)	Best-of-K (K=8)	S3 (N=4, b=2)	提升 (vs Baseline)
GSM8K (数学)	68.16%	69.56%	70.21%	+2.05 pp
MATH-500 (高阶数学)	25.60%	28.20%	30.20%	+4.60 pp
TruthfulQA (事实性)	46.49%	49.36%	49.57%	+3.08 pp
ARC-Challenge (推理)	76.11%	79.30%	77.86%	+1.75 pp

关键发现：
- 数学任务提升最大： 在 MATH-500 上，S3 比 Best-of-K 高出 2.0 个百分点，表明在需要多步推理的任务中，去噪过程中的轨迹搜索比单纯增加最终采样数更有效。
- 计算效率： S3 在相同的计算预算（NFE, Number of Function Evaluations）下，性能优于 Best-of-K 的帕累托前沿。
- 块长度敏感性： S3 在细粒度块长度（Block Length $K \in \{2, 4, 16\}$ ）下表现最佳，因为此时前瞻信号更准确；在粗粒度块长度下，Best-of-K 在 ARC-Challenge 上偶尔表现更好。
- 消融实验： 单独使用“前瞻搜索”或“奖励倾斜”均不足以达到最佳效果，两者的结合（S3）产生了协同效应。

5. 意义与结论 (Significance & Conclusion)

范式转变： 本文证明了对于扩散语言模型，测试时扩展（Test-time Scaling）的关键在于在去噪过程中重新分配计算资源，而不仅仅是增加最终输出的采样数量。
无需重训练： S3 是一种即插即用的推理时策略，不需要修改模型权重或训练新的奖励模型（仅需轻量级规则验证器），具有极高的实用价值。
通用性： 该方法不仅适用于数学推理，也提升了事实性问答和常识推理任务的表现，揭示了 DLMs 去噪轨迹搜索的通用潜力。
局限性： 方法依赖于验证器的质量。如果验证器信号噪声大或与任务不匹配，可能会误导轨迹搜索。此外，粒子扩展和重采样增加了推理延迟，需要在计算成本与性能之间进行权衡。

总结： S3 通过引入经典的粒子滤波思想（结合前瞻评分和重采样），成功解决了扩散语言模型中分布错位的问题，为 DLMs 的测试时扩展提供了一条切实可行的新路径。

S3S^3S3: Stratified Scaling Search for Test-Time in Diffusion Language Models