掩码扩散模型中的掩码调度

在扩散模型(Diffusion Models)尤其是掩码扩散模型(Mask Diffusion Models, MDM)的研究中,Noise Scheduler 的设计是决定模型学习上限与采样质量的核心枢纽。本文对 MDM 中的 Noise Scheduler 进行了全面而深入的统一分析,系统性地探讨了当前主流的调度策略,通过对比分析各策略的优缺点,本文揭示了精细化调度策略如何重塑 Token 间的依赖动力学,并在最后对未来突破底层逻辑缺陷的发展方向进行了展望。

Introduction

在扩散模型(Diffusion Models)的研究中,Noise Scheduler 的选取不仅决定了采样的轨迹,更直接关系到模型对复杂分布的学习上限。相比于连续域的高斯扩散,掩码扩散模型(Mask Diffusion Models, MDM) 在离散数据建模中展现了独特的潜力。然而,MDM 普遍采用的 Absorb Transition Kernel 在赋予模型简洁逻辑的同时,也引入了三个关键的理论与工程挑战:

首先,曝光偏差(Exposure Bias)在 MDM 中尤为严峻。由于 Absorb 机制导致生成的 Token 具有“不可修改”的特性,采样过程中的任何微小偏离都会因缺乏纠错机制而随时间步迅速累积,导致生成序列的崩塌。

其次,数据内在序(Intrinsic Order)的复杂性限制了训练效率。现实数据(如自然语言或代码)往往存在多样的依赖顺序,而 MDM 试图在有限的参数空间内捕捉所有可能的生成路径。由于无法在多项式时间内穷举并对每一种生成顺序都实现最优学习,模型往往面临训练效率不足的问题,难以在复杂的结构化数据上达到收敛。 最根本的挑战在于 独立性假设带来的概率偏离。MDM 在建模阶段往往假设各 Token 之间是条件独立的,这与离散数据强耦合的本质相悖,导致模型估计的边缘分布乘积严重偏离真实的联合概率分布。

因此,设计精密的 Noise Scheduler 不再仅仅是掩码比例的简单伸缩,而是弥补 MDM 底层逻辑缺陷的关键枢纽。一个优化的调度器不仅需要通过精细的掩码率平衡来缓解曝光偏差,更核心的在于通过时间轨迹的演变来重塑 Token 间的依赖动力学。

除了寻求 “零全相关路径” 以在统计上近似独立分布外,更具前瞻性的修正方案在于打破单一 Token 的处理范式。例如,通过引入Block-wise Masking,Scheduler 可以强制模型捕捉局部强耦合的语义簇,从物理层面缓解独立性假设带来的联合概率偏离。同时,引入Iterative Remasking 赋予了模型在推理轨迹上的“反悔”与修正能力,打破了 Absorb Kernel 的单向不可逆性。这些策略通过 Scheduler 的动态协调,不仅在统计意义上降低了信息熵损失,更在本质上充当了独立性假设与结构化依赖之间的校准器,从而在保证计算效率的前提下,实现对离散数据深层联合概率分布的精准还原。

基于以上观察,本文调研并总结了当前 Mask Diffusion Model 领域针对Noise Scheduler的改进方案,重点探讨如何通过调度的精细化设计来突破上述瓶颈。

Background

Masked diffusion model

Forward Process

设token共有 $C$ 种,加上mask,则每个token的状态空间为${0,1,…,C}$。设token sequence表示为 $x=x^1x^2…x^L$,$x_0\sim p_{data}$,那么条件概率 \(q(x_t\mid x_0)=\prod_{i=1}^Lq(x_t^i\mid x_0^i)\),其中 $q(x_t^i\mid x_0^i)=(1-\alpha_t)\delta_m(x_t^i)+\alpha_t\delta_{x_0^i}(x_t^i)$ 表示token在 $t$ 时刻,以$1-α_t$ 概率变为mask,以 $α_t$ 概率保持不变,且各token的变化是独立的。

Reverse Precess

设 $ 0\leq s\le t\leq 1$,在反向过程中,为了计算 $p(x_s^i\mid x_t^i)$,我们首先计算 $p(x_t^i\mid x_s^i)$:

\[q(x_t^i \mid x_s^i) = \begin{cases} \frac{\alpha_s - \alpha_t}{\alpha_s} \delta_{\text{m}}(x_t^i) + \frac{\alpha_t}{\alpha_s} \delta_{x_s^i}(x_t^i) & \text{if } x_s^i \in \mathcal{X}, \\ \delta_{x_s^i}(x_t^i) & \text{if } x_s^i = \text{m}. \end{cases}\]

并由贝叶斯公式得:

\[q(x_s^i \mid x_t^i, x_0^i) = \begin{cases} \delta_{x_t^i}(x_s^i) & \text{if } x_t^i \in \mathcal{X}, \\ \frac{1 - \alpha_s}{1 - \alpha_t} \delta_{\text{m}}(x_s^i) + \frac{\alpha_s - \alpha_t}{1 - \alpha_t} \delta_{x_0^i}(x_s^i) & \text{if } x_t^i = \text{m}. \end{cases}\]

这表明,当token $x_t^i$ 是mask时,以 $\frac{1 - \alpha_s}{1 - \alpha_t} $的概率保持为 mask,以 $ \frac{\alpha_s - \alpha_t}{1 - \alpha_t} $ 的概率变为 $x_0^i$,而当token unmask后,则不再参与更新。

Model and Loss

由于 $p(x_s\mid x_t)=\mathbb E_{p(x_0\mid x_t)}[q(x_s\mid x_t,x_0)]$,反向过程可以通过先采样 $x_0\sim p(\cdot\mid x_t)$,然后再采样 $x_s\sim q(\cdot\mid x_t,x_0)$来实现。由于 $x_0$ 是未知的,通过 $p_{\theta}(x_0\mid x_t)=\prod_{i=1}^Lp_{\theta}(x_0^i\mid x_t)$来建模,并以如下损失函数来优化:

\[\mathcal{L}_{\text{vb}}(\mathbf{x}_0; \theta) = \int_0^1 \frac{\alpha_t'}{1 - \alpha_t} \mathbb{E}_{q(\mathbf{x}_t \mid \mathbf{x}_0)} \left[ \sum_{i=1}^{L} \log p_\theta(x_0^i \mid x_t^i) \right] dt,\]

Inference

在推理阶段,给定一个完全被掩码的序列 $\mathbf{x_L} = \left[\text{m}, \text{m}, \dots, \text{m}\right]$ ,我们需要通过 $T$ 步去噪得到 $\mathbf{x_0}$。根据之前的讨论,每一时刻 $t$ 的状态转移不仅取决于模型的预测能力 $p_\theta(x_0^i \mid \mathbf{x_t})$,更取决于采样路径(Sampling Path)的选择。

我们将时间 $t+1\rightarrow t$ 的推理建模为两步,首先根据某种策略 $\pi(i\mid x_{t+1}),i=1,2,\cdots L$ 选出被unmask的位置,记为 $M_{t+1}$,然后对 $M_{t+1}$ 中的位置,按照模型预测的logist进行采样,即 $x^i_t\sim p_{\theta}(x_t^i\mid x_{t+1})$。那么一次推理即为:

\[x_t^i = \begin{cases} \sim p_{\theta}(x^i \mid \mathbf{x}_{t+1}), & \text{if } i \in M_{t+1} \quad \text{(在当前步被选中Unmask)} \\ x_{t+1}^i, & \text{if } i \notin M_{t+1} \quad \text{(保持上一时刻状态)} \end{cases}\]

也即:

\[[\mathbf{T}_{t+1 \rightarrow t}^M](\mathbf{x}_{t+1}, \mathbf{x}_t) = \begin{cases} \underbrace{\pi(M \mid \mathbf{x}_{t+1})}_{\text{Selection}} \cdot \underbrace{\prod_{i \in M} p_\theta(x_t^i \mid \mathbf{x}_{t+1})}_{\text{Generation}} & \text{if } \mathbf{x}_t^{-M} = \mathbf{x}_{t+1}^{-M} \\ 0 & \text{otherwise} \end{cases}\]

Exposure Bias

首先,我们定义两个核心概率测度,它们存在于相同的状态空间 $\mathcal{X}^L$ 上:

\[P_{t}(\mathbf{x_t}) = \int p_{\text{data}}(\mathbf{x_0}) q(\mathbf{x_t} \mid \mathbf{x_0}) d\mathbf{x_0}\] \[\hat{P}_{t}(\hat{\mathbf{x}}_t) = \sum_{\hat{\mathbf{x}}_{t+1}} \hat{P}_{t+1}(\hat{\mathbf{x}}_{t+1}) \mathbf{T}_{t+1 \to t}^M(\hat{\mathbf{x}}_t \mid \hat{\mathbf{x}}_{t+1})\]

我们关注的目标是时刻 $t$ 的总分布偏移:$\Delta_t = D_{KL}(P_t \Vert \hat{P_t})$。

利用 KL 散度的链式法则,我们可以建立 $\Delta_t$ 的演化关系。

设 $\mathcal{Q}_{t+1 \to t}$ 为真实数据分布诱导的最优反向转移算子。显然有 $P_t = \mathcal{Q}_{t+1 \to t} P_{t+1}$。则:

\[\Delta_t = D_{KL}\left( \mathcal{Q}_{t+1 \to t} P_{t+1} \big\| \mathbf{T}_{t+1 \to t}^M \hat{P}_{t+1} \right)\]

通过引入中间项 $\mathbf{T}_{t+1 \to t}^M P_{t+1}$ 并利用散度的三角形性质,可推导出如下上界:

\[\Delta_t \le \underbrace{D_{KL}(\mathcal{Q}_{t+1 \to t} P_{t+1} \| \mathbf{T}_{t+1 \to t}^M P_{t+1})}_{\text{单步局部误差 } \mathcal{E}_{\text{loc}}} + \underbrace{D_{KL}(\mathbf{T}_{t+1 \to t}^M P_{t+1} \| \mathbf{T}_{t+1 \to t}^M \hat{P}_{t+1})}_{\text{历史误差传播}}\leq \mathcal{E}_{\text{loc}}+\Delta_{t+1}\]

通过优化 $\mathbf{T}_{t+1 \to t}^M$ ,我们可以同时减小单步局部误差从而减少 $\Delta_t$。

Intrinsic Order

我们首先形式化“数据内在序”这一概念。设完整序列为

\[\mathbf{x}_0 = (x_0^1, x_0^2, \dots, x_0^L) \sim p_{\text{data}}\]

对于任意子集 $S \subset {1,\dots,L}$,定义其条件生成难度为:

\[\mathcal{H}(S) \;\triangleq\; \mathbb{E}_{\mathbf{x}_0 \sim p_{\text{data}}} \left[ H\big( \mathbf{x}_0^S \mid \mathbf{x}_0^{\bar S} \big) \right]\]

该量刻画:在其余 token 已知的条件下,集合 $S$ 中 token 的不确定性。

于是数据分布 $p_{\text{data}}$ 诱导了一族等价的最优生成序:

\[\mathcal{O}^* = \arg\min_{\pi \in \mathfrak{S}_L} \sum_{k=1}^L \mathcal{H}\big(\{\pi_k\} \mid \{\pi_1,\dots,\pi_{k-1}\}\big)\]

其中 $\pi$ 为 token 的一个排列(生成顺序)。

但在分布层面定义单一或固定的最优生成顺序在实践中是不合理的,不同样本 $\boldsymbol{x}_0$ 往往对应不同的 $\pi^*(\boldsymbol{x}_0)$。这一现象反映了真实离散数据中普遍存在的生成顺序多样性。对于具体的样本 $x_0\sim p_{data}$,其最优生成顺序一般依赖于所体现的局部结构与依赖关系,可以形式化为:

\[\pi^*(\boldsymbol{x}_0)=\operatorname*{arg\,min}_{\pi\in \mathfrak{S}_L }\sum_{k=1}^LH\bigg(x_0^{\pi_k}\mid \boldsymbol{x}_0^{\{\pi_1,\cdots,\pi_{k-1}\}}\bigg)\]

Impact on Training Complexity

在标准的 MDM 训练中,噪声调度器通常采用均匀采样策略,即假设所有 Token 被掩码的概率是均等的。这意味着模型被迫在训练过程中学习所有可能的生成排列 $\pi \in \mathfrak{S}_L$ 的平均。

然而,对于具有强结构依赖的数据,不同生成路径的学习难度存在巨大差异。我们可以将训练过程中的期望损失分解为基于排列的加权和:

\[\mathcal{L}_{\text{total}} \approx \mathbb{E}_{\mathbf{x}_0} \mathbb{E}_{\pi \sim U(\mathfrak{S}_L)} \left[ \sum_{k=1}^L -\log p_\theta(x_0^{\pi_k} \mid \mathbf{x}_0^{\{\pi_1, \dots, \pi_{k-1}\}}) \right]\]

由于参数空间的限制,模型无法在多项式时间内拟合所有 $L!$ 种路径的边缘分布。均匀调度强迫模型在早期时间步去预测那些高熵、高依赖性的 Token,这导致了优化目标中的梯度冲突。

这种冲突不仅拖慢了收敛速度,更在理论上降低了模型所能达到的性能上界。我们可以定义Optimization Gap,即 $\Delta \mathcal{L}_{\text{order}}$,来量化由于忽略数据内在序而引入的额外复杂性:

\[\begin{aligned} \Delta \mathcal{L}_{\text{order}} = \mathbb{E}_{\mathbf{x}_0} \left[ \sum_{k=1}^L \left( \underbrace{\mathbb{E}_{\pi \sim U} [H(x_0^{\pi_k} \mid \mathbf{x}_0^{\pi_{\lt k}})]}_{\text{Avg. Complexity (Uniform)}} - \underbrace{H(x_0^{\pi^*_k} \mid \mathbf{x}_0^{\pi^*_{\lt k}})}_{\text{Min. Complexity (Optimal)}} \right) \right] \end{aligned}\]

其中 $\pi^*$ 代表使得序列条件熵和最小的最优生成序。$\Delta \mathcal{L}_{\text{order}}$ 的存在揭示了一个核心矛盾:模型不仅要学习符合因果律的“易学路径”($P(\text{Result}\mid\text{Cause})$),还被迫消耗大量参数去拟合违反直觉的逆向或乱序路径 $P(\text{Cause}\mid\text{Result})$。

这一现象引发了容量稀释问题。根据 Rate-Distortion Theory,如果我们将模型参数 $\theta$ 视为容量有限的信息信道 $C_\theta$,那么均匀采样策略人为地将任务所需的信息编码率从 $R_{\text{opt}}$ 提升到了 $R_{\text{uni}}$。当 $R_{\text{uni}} > C_\theta$ 时,模型无法完美拟合这一高熵混合分布,必然产生不可约减的失真。

换言之,$\Delta \mathcal{L}_{\text{order}}$ 实质上度量了有多少模型容量被浪费在了拟合无效的噪音路径上。这种资源的被动分散,导致模型在处理真正重要的正向依赖时可用参数减少,从而迫使模型学到的联合分布 $P_\theta(\mathbf{x}_0)$ 呈现出一种平滑化的特征,难以捕捉数据分布中原本尖锐的多模态结构。因此,设计能逼近 $\pi^*$ 的 Scheduler,本质上是在进行一种课程学习,旨在消除 $\Delta \mathcal{L}_{\text{order}}$ 带来的熵增,从而在有限的参数规模下逼近更高的性能上限。

Impact on Inference Accuracy

在推理阶段,当噪声调度器 $\pi$ 选取的去噪路径违背数据的内在序 $\pi^*$ 时,会引发级联式的精度崩塌。

当调度器强迫模型在缺乏Context的情况下优先预测下游 Token 时,模型面临的是一个高熵分布。

设 $x^A \to x^B$ 为真实因果依赖,若 $\pi$ 优先揭露 $x^B$,则模型需从边缘分布采样:

\[x^B \sim p_\theta(x^B \mid \emptyset) = \sum_{x^A} p(x^B \mid x^A)p(x^A)\]

相比于条件分布 $p(x^B\mid x^A)$,该分布极其平坦,导致采样极易命中低概率或错误的token。

Token Independence

让 $\mathcal{M}_t$ 表示 $t$ 时刻被 Mask 的集合,$c = x_{t+1}$ 为当前上下文。我们需要在 $t+1 \to t$ 这一步解开一个子集 $\mathcal{U} \subset \mathcal{M}_{t+1}$。模型使用独立分布 $Q_\theta(x_\mathcal{U} \mid c) = \prod_{i \in \mathcal{U}} Q_\theta(x_i \mid c)$ 来近似真实分布 $P(x_\mathcal{U} \mid c)$。我们要最小化的目标是 $D_{\text{KL}}(P \Vert Q_\theta)$。通过引入真实的边缘分布乘积 $\prod_{i \in \mathcal{U}} P(x_i \mid c)$ 作为中间项,我们可以将这个 KL 散度完美分解:

\[\begin{aligned} D_{\text{KL}}(P \| Q_\theta) &= \mathbb{E}_{x \sim P} \left[ \log \frac{P(x_\mathcal{U}\mid c)}{\prod_{i \in \mathcal{U}} Q_\theta(x_i\mid c)} \right] \\ &= \mathbb{E}_{x \sim P} \left[ \log \left( \frac{P(x_\mathcal{U}\mid c)}{\color{blue}{\prod_{i \in \mathcal{U}} P(x_i\mid c)}} \cdot \frac{\color{blue}{\prod_{i \in \mathcal{U}} P(x_i\mid c)}}{\prod_{i \in \mathcal{U}} Q_\theta(x_i\mid c)} \right) \right] \\ &= \underbrace{\mathbb{E}_{x \sim P} \left[ \log \frac{\prod_{i \in \mathcal{U}} P(x_i\mid c)}{\prod_{i \in \mathcal{U}} Q_\theta(x_i\mid c)} \right]}_{\text{Term 1: Sum of Marginal Errors}} + \underbrace{\mathbb{E}_{x \sim P} \left[ \log \frac{P(x_\mathcal{U}\mid c)}{\prod_{i \in \mathcal{U}} P(x_i\mid c)} \right]}_{\text{Term 2: Total Correlation (Dependency)}} \end{aligned}\]

基于上述 Total Correlation (TC) 的数学分解,我们可以深刻理解 Noise Scheduler 在推理阶段的核心使命。由于模型参数 $\theta$ 在推理时是固定的,Term 1 (Marginal Errors) 的上限在训练结束时即已确定。因此,推理质量的提升完全取决于能否通过 Scheduler 的策略性选择,将 Term 2 (Total Correlation) 最小化。

换言之,在推理过程中,Scheduler 不再仅仅是一个简单的进度条控制者(决定 Unmask 多少个 Token),而是扮演了 “依赖解耦器” 的角色。

Remark

回顾前文关于 Exposure Bias 的推导,总误差 $\Delta_t$ 的累积源于每一步的局部误差 $\mathcal{E}_{\text{loc}}$。在推理的 $t+1 \to t$ 阶段,给定上下文 $c = \mathbf{x}_{t+1}$,Scheduler $\pi$ 首先决定unmask的子集 $\mathcal{U} \sim \pi(\cdot\mid c)$,随后模型 $p_\theta$ 对其进行填充。 我们将单步局部误差 $\mathcal{E}_{\text{loc}}$ 重新形式化为关于策略 $\pi$ 的期望形式:

\[\mathcal{E}_{\text{loc}}(\pi, \theta) = \mathbb{E}_{\mathcal{U} \sim \pi(\cdot \mid c)} \left[ D_{KL}\Big(~P(x_\mathcal{U} \mid c) ~\Big\| \underbrace{\prod_{i \in \mathcal{U}} p_\theta(x_i \mid c)}_{\text{Independence Assumption}} \Big) \right]\]

我们可以将这一误差进一步拆解为两部分,从而清晰地分离 Intrinsic Order 和 Independence 的影响:

\[\mathcal{E}_{\text{loc}}(\pi, \theta) = \mathbb{E}_{\mathcal{U} \sim \pi} \left[ \underbrace{\sum_{i \in \mathcal{U}} D_{KL}(P(x_i\mid c) \| p_\theta(x_i\mid c))}_{\text{(I) Marginal Precision Error}} + \underbrace{\mathbb{I}(\lvert\mathcal{U}\rvert > 1) \cdot \text{TC}(x_\mathcal{U} \mid c)}_{\text{(II) Structural Dependency Error}} \right]\]

其中 $\mathbb{I}(\cdot)$ 为示性函数,强调项 (II) 仅在一步解开多个 Token 时存在。

总结而言,数据的内在序以及独立性同时影响了训练与推理,当单步解码数 $\mathcal U$ 较大时,独立性会显著制约模型精度,但同时也降低了对数据内在序的违背;而当单步解码数 $\mathcal U$ 很小时,独立性的影响大大减弱,但数据内在序的制约增强。通过设计合理的mask/unmask策略 $\pi$,使得模型训练和推理能够顺应数据的内在序,并尽可能缓解Independence 的影响,我们可以大大提高模型的训练效率、模型容量并缓解exposure bias。

Current Schedulers

Random Schedulers

随机策略是 MDM 的基准实现(MDLM, MD4, SEDD),它假设所有 token 的地位是平等的,忽略了数据的内在序(Intrinsic Order)。

Confidence-based Schedulers

与随机策略不同,基于置信度(Confidence-based)的策略 (Train for the worst, Plan for the best, MGDM ) 试图通过模型的预测结果动态地寻找“易学路径”。其核心逻辑是:模型越有把握(置信度越高)的 token,越倾向于先被解开(Unmask),以此顺应数据的内在序 $\pi^*$。在该框架下,策略 $\pi$ 不再是均匀分布,而是取决于模型输出的概率分布 $p_\theta(x^i \mid \mathbf{x}_{t+1})$。

图示比较了不同 Sampling Steps 下的生成困惑度(左轴)与预测熵(右轴):置信度驱动的自适应推理在较少步数时即可保持更低困惑度,并以更平滑的熵轨迹逼近高步数性能。

Planer-based Schedulers

这一部分我们介绍一些通过模型学习得到的scheduler,从而减少归纳偏置,得到更好的泛化性。

Noise Predict Planer

Paper: THINK WHILE YOU GENERATE: DISCRETE DIFFUSION WITH PLANNED DENOISING (ICLR 2025)

DDPD 由两个独立的神经网络组成:

DDPD(Discrete Diffusion with Planned Denoising)框架示意图:模型由 Planner 与 Denoiser 两部分组成。Planner 先对每个位置输出噪声状态概率 $p_\theta(z_t^d\mid x_t)$,识别当前最需要修复的 token;随后 Denoiser 在被选中的噪声位置上进行条件预测 $p_{1\mid t}^\theta(x_1^d\mid x_t,z_t^d=N)$。与固定步长调度不同,DDPD 通过估计全局噪声比例反推实际时间嵌入 $\tilde t=\alpha_t^{-1}\!\left(1-\frac{1}{D}\sum_{d'}p_\theta(z_t^{d'}=N\mid x_t)\right)$,从而按“先规划、后去噪”的顺序迭代生成,提高采样效率与最终质量。

RL-Policy Planer

RL-Policy Planer 旨在将离散扩散模型的去噪过程建模为一个马尔可夫决策过程,通过强化学习来学习最优的去噪轨迹或 Unmask 策略。与基于启发式的调度器(如随机或置信度采样)不同,RL-Policy 能够根据下游任务的最终奖励(如数学题的准确率、推理的逻辑正确性)自适应地调整生成路径,从而在更复杂的生成场景中展现出优越性。

Freeze backbone with trainable Policy head

Paper: IMPROVING DISCRETE DIFFUSION UNMASKING POLICIES BEYOND EXPLICIT REFERENCE POLICIES (ICLR 2026)

这种架构的核心思想是分工明确:底层的 Backbone已经具备了强大的条件概率建模能力,而 Policy Head 则作为一个scheduler,专门负责根据当前上下文评估各位置的解掩码优先级。

冻结 Backbone、仅训练 Policy Head 的 RL-Policy Planner 架构示意:在每一步 $t+1\to t$,策略网络基于当前序列状态输出各位置的解掩码优先级 $\pi_\phi(i\mid x_{t+1})$,选择需要优先去噪的 token,再由冻结的基础 MDM 执行条件生成。该设计将“位置决策”与“内容生成”解耦,在保持训练开销较低的同时,通过奖励信号优化采样顺序以提升推理质量。

Trainable backbone with fixed path

Paper: MDPO: OVERCOMING THE TRAINING-INFERENCE DIVIDE OF MASKED DIFFUSION LANGUAGE MODELS

MDPO generates a group of answers given a query for RL rollouts. Then all completions at intermediate and final steps are verified with a reward model. Based on verified rewards, MDPO estimates the advantage of step t by considering rewards of the other steps in the current rollout and step t from other rollouts in the group. These estimated advantages are used for policy optimization.

与上述方案不同,MDPO 并不改变位置选择的逻辑(依然沿用 Confidence-based 启发式),而是通过 RL 训练直接优化 Backbone 自身,使其输出的 Logit 能够更完美地适配推理时的调度轨迹。

Manual Selected Schedulers

我们在这一部分介绍一些人为选定的scheduler

Subsequence-based Schedulers (SPMDM)

Paper: SPMDM: Enhancing Masked Diffusion Models through Simplifying Sampling Path

与传统的全局随机调度不同,SPMDM 认为高效的采样路径应具备两个特征:局部优先(优先解开已解 token 附近的掩码)和逻辑顺序优先(优先解开推理链条中靠前的子序列)。

SPMDM partitions the input sequence x into K subsequences and introduces noise with varying magnitudes to each of them. Compared to MDLM, our method not only models token-level dependencies but also explicitly encourages the learning of inter-subsequence relationships. In contrast to BDLM, SP

Block Diffusion Language Model

Paper: BLOCK DIFFUSION: INTERPOLATING BETWEEN AUTOREGRESSIVE AND DIFFUSION LANGUAGE MODELS

BD3-LMs 提出了一种在自回归(AR)建模与离散扩散建模之间进行插值的新范式。其核心动机在于通过引入块级(Block-wise)的结构,克服标准扩散模型在生成任意长度序列、推理效率(KV Caching)以及似然建模质量(Perplexity)上的局限性。

Block diffusion sequentially generates blocks of tokens by performing diffusion within each block and conditioning on previous blocks. By combining strength from autoregressive and diffusion models, block diffusion overcomes the limitations of both approaches by supporting variable-length, higher-quality generation and improving inference efficiency with KV caching and parallel sampling.

BD3-LM 将长度为 $L$ 的序列划分为 $B$ 个长度为 $L’$ 的块。模型在块之间遵循自回归分解,而在每个块内部执行离散扩散过程:

\[\log p_{\theta}(x) = \sum_{b=1}^{B} \log p_{\theta}(x^{b} \mid x^{\lt b})\]

Final Remark

Schedulers Policy $\pi(M \mid c)$ Sampler $p_{\theta}$ Advantage Shortage
Random Schedulers $\pi(\mathcal{U}) = \displaystyle\frac{1}{\binom{ \vert \mathcal{M}_{t+1}\vert}{k}}$ 所有可能的置换路径 $\mathfrak{S}_L$ 的平均分布 前向过程有闭式解;
无需在采样过程中进行额外的 Logit 排序或规划器前向计算
训练资源浪费;
无视内在序,误差累积严重
Confidence-based Schedulers $\mathcal{U} = \text{arg-topk}_{i \in \mathcal{M}_{t+1}} { s_i \mid s_i = \max_{x \in \mathcal{X}} \log p_\theta(x^i = x \mid \mathbf{x}_{t+1}) }$
$\pi(\mathcal{U} \mid \mathbf{x}_{t+1}) = \mathbb{1}(\mathcal{U} = \text{arg-topk})$
通常在训练时引入 Token-level 重加权,使模型对低置信度(难样本)的边界刻画更清晰。 能根据不同的输入样本动态调整生成顺序 模型训练不充分或词表过大,Logit 排序是嘈杂的;
无法解决 Total Correlation 误差
Planer-based Schedulers $\mathcal{U} = { i \mid p_\phi(z^i = N \mid \mathbf{x}_{t+1}) < \tau }$ $p_\theta$ 只负责被 $p_\phi$ 选中的维度,其输入通常包含由规划器估计的实时时间步 $\tilde{t}$。 支持 Remasking;
避免容量稀释;
逐个 token 生成,且有 remask,效率低下;
RL-Policy Schedulers $\pi_\phi(i \mid \mathbf{x}_{t+1}) = \text{Softmax}(h_\phi(\mathbf{x}_{t+1}))$ 训练目标从简单的最大似然转向对高奖励路径的偏好 训推一致;
性能上限高
训练复杂
Manual/Block Schedulers $\pi(\mathcal{U}) \implies \mathcal{U} \subset \text{Block}_k,\ \text{其中 } k \text{ 由预设顺序或块噪声 } t_k \text{ 决定}$ 学习局部强耦合语义 显著降低了 $\text{TC}(x_\mathcal{U} \mid c)$ 固定的块划分可能无法完全契合数据的非线性逻辑依赖



从统一视角看,Mask Diffusion Model 中的 Noise Scheduler 本质上是一个“结构约束器(Structural Regulator)”,而非传统意义上的时间控制器。它通过决定 何时、解开什么、以及 一次解多少 Token,直接塑造了模型在训练与推理中所面对的条件分布,从而决定了三类核心误差的权衡关系:

  1. Exposure Bias 的时间传播误差(由单步局部偏差累积而来);
  2. Intrinsic Order 被违背所引入的优化复杂度与容量稀释;
  3. Token Independence 假设下不可避免的 Total Correlation 偏离。

分析表明,这三者并非独立问题,而是通过 Scheduler 的策略性选择在同一条误差曲线上被共同调节:

综合来看,RL-Policy Scheduler 通过reward,同时兼顾了内在序与全相关误差的制约,并且其采样路径始终是遵循policy的,更加训推一致,生成效果趋于最好。

Enjoy Reading This Article?

Here are some more articles you might like to read next:

  • Masking Schedulers of Mask Diffusion Model
  • Flow Matching and Continuous Normalizing Flows
  • The Unification of DDPM and Score-based Models