在扩散模型(Diffusion Models)尤其是掩码扩散模型(Mask Diffusion Models, MDM)的研究中,Noise Scheduler 的设计是决定模型学习上限与采样质量的核心枢纽。本文对 MDM 中的 Noise Scheduler 进行了全面而深入的统一分析,系统性地探讨了当前主流的调度策略,通过对比分析各策略的优缺点,本文揭示了精细化调度策略如何重塑 Token 间的依赖动力学,并在最后对未来突破底层逻辑缺陷的发展方向进行了展望。
在扩散模型(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的改进方案,重点探讨如何通过调度的精细化设计来突破上述瓶颈。
设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的变化是独立的。
设 $ 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后,则不再参与更新。
由于 $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,\]在推理阶段,给定一个完全被掩码的序列 $\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}\]首先,我们定义两个核心概率测度,它们存在于相同的状态空间 $\mathcal{X}^L$ 上:
我们关注的目标是时刻 $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$。
我们首先形式化“数据内在序”这一概念。设完整序列为
\[\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)\]在标准的 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}}$ 带来的熵增,从而在有限的参数规模下逼近更高的性能上限。
在推理阶段,当噪声调度器 $\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。
让 $\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}\]Term 1: 边缘误差之和 (对应标准 CE Loss)
\[\text{Term 1} = \sum_{i \in \mathcal{U}} D_{\text{KL}}(P(x_i\mid c) \| Q_\theta(x_i\mid c))\]由于 $D_{\text{KL}}(P_i | Q_i) = \text{CrossEntropy}(P_i, Q_i) - H(P_i)$,且 $H(P_i)$ 对于模型参数 $\theta$ 是常数,最小化这一项完全等价于最小化标准 Cross Entropy Loss。
Term 2: 全相关误差
\[\text{Term 2} = D_{\text{KL}}\left( P(x_\mathcal{U}\mid c) \bigg\| \prod_{i \in \mathcal{U}} P(x_i\mid c) \right) = \text{TC}(x_\mathcal{U}\mid c)=\left( \sum_{i \in \mathcal{U}} H(x_i\mid c) \right) - H(x_\mathcal{U}\mid c)\]这衡量了变量之间偏离独立性的程度。如果 $x_\mathcal{U}$ 中的 token 在给定 $c$ 下是完全条件独立的,则此项为 0。依赖关系越强,此项越大。
基于上述 Total Correlation (TC) 的数学分解,我们可以深刻理解 Noise Scheduler 在推理阶段的核心使命。由于模型参数 $\theta$ 在推理时是固定的,Term 1 (Marginal Errors) 的上限在训练结束时即已确定。因此,推理质量的提升完全取决于能否通过 Scheduler 的策略性选择,将 Term 2 (Total Correlation) 最小化。
换言之,在推理过程中,Scheduler 不再仅仅是一个简单的进度条控制者(决定 Unmask 多少个 Token),而是扮演了 “依赖解耦器” 的角色。
回顾前文关于 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。
随机策略是 MDM 的基准实现(MDLM, MD4, SEDD)
Training 在训练时,时间步 $t \in [0, 1]$ 均匀采样,对应的 mask 比例为 $1 - \alpha_t$。对于序列中的任意位置 $i $,其被选中的概率是独立的:
\[\pi(M \mid \mathbf{x}_0) = \binom{L}{k}^{-1}, \quad \text{where } k \approx (1-\alpha_t)L\]此时,损失函数 $\mathcal{L}_{\text{vb}}$ 退化为对所有可能排列 $\mathfrak{S}_L$ 的均匀期望。
Inference 在每一推理步 $t+1 \to t$,策略 $\pi$ 随机从当前 mask 集合中挑出固定比例(由 scheduler 决定)的位置进行 unmask:
\[M_{t+1} \subset \{i \mid x_{t+1}^i = \text{m}\}, \quad \text{s.t. } \lvert M_{t+1} \rvert = \lfloor (\alpha_t - \alpha_{s})L \rfloor\]选择过程与上下文 $c$ 无关,即 $\pi(M \mid \mathbf{x}_{t+1}) = \pi(M)$。模型在每一步独立预测所有被掩码位置的概率分布:
\[x_t^i \sim p_\theta(x^i \mid \mathbf{x}_{t+1}) \quad \forall i \in M_{t+1}\]与随机策略不同,基于置信度(Confidence-based)的策略 (Train for the worst, Plan for the best, MGDM )
Inference ($\pi_{\text{inf}}$):
在 $t+1 \to t$ 步,计算所有 mask 位置的 $s_i$,并选择得分最高的前 $k$ 个位置:
\[M_{t+1} = \text{arg-topk}_{i} \{s_i\}, \quad \text{where } k = \lfloor (\alpha_t - \alpha_s)L \rfloor\]对应的生成过程为:
\[\pi(M \mid \mathbf{x}\_{t+1}) = \begin{cases} 1 & \text{if } M = M_{t+1} \\ 0 & \text{otherwise} \end{cases}\]为了保留一定的随机性,也可以设计 topk ratio 来控制topk选择的比例,即将一部分的topk和一部分随机选出的token放入 $M_{t+1}$ 中。
Training
当前的工作中,并没有在训练中显式实现使用confidence-based Policy进行mask的,当有一些工作通过对损失函数的重加权来实现对特定token或者说对特定生成序的偏好训练,如在 MGDM
其中,$v(\mathbf{x}_{t,n}) = \alpha(1 - \exp(-u(\cdot)))^\beta$ 是自适应令牌级重加权项。当 $\beta > 0$ 时,该项会减小易学 Token(低 Loss)的权重,并放大难学 Token(高 Loss)的权重。从优化目标上看,这种训练策略旨在通过“强迫”模型关注那些当前预测不准确的难样本,从而提高整体训练效率。
Advantage:
Shrotage:
这一部分我们介绍一些通过模型学习得到的scheduler,从而减少归纳偏置,得到更好的泛化性。
Paper: THINK WHILE YOU GENERATE: DISCRETE DIFFUSION WITH PLANNED DENOISING (ICLR 2025)
DDPD 由两个独立的神经网络组成:
规划器 ($p_\theta(z_t^d \mid x_t)$):执行二分类任务,预测序列中每个维度 $d$ 是否被损坏(即处于噪声状态 $z_t^d = N$)。它仅为每个维度输出 1 个 Logit。
去噪器 ($p_{1\mid t}^\theta(x_1^d \mid x_t, z_t^d = N)$):在给定位置被判定为噪声的条件下,预测其真实的数据分布。 DDPD 不再依赖预设的时间步长,而是通过规划器估计的全局噪声水平来反推实际时间: $\tilde{t} = \alpha_t^{-1}(1 - \sum_{d’} p_\theta(z_t^{d’} = N \mid x_t) / D)$ 并将实际时间作为denoiser的时间嵌入。
Advantage:
Shortage:
RL-Policy Planer 旨在将离散扩散模型的去噪过程建模为一个马尔可夫决策过程,通过强化学习来学习最优的去噪轨迹或 Unmask 策略。与基于启发式的调度器(如随机或置信度采样)不同,RL-Policy 能够根据下游任务的最终奖励(如数学题的准确率、推理的逻辑正确性)自适应地调整生成路径,从而在更复杂的生成场景中展现出优越性。
Paper: IMPROVING DISCRETE DIFFUSION UNMASKING POLICIES BEYOND EXPLICIT REFERENCE POLICIES (ICLR 2026)
这种架构的核心思想是分工明确:底层的 Backbone已经具备了强大的条件概率建模能力,而 Policy Head 则作为一个scheduler,专门负责根据当前上下文评估各位置的解掩码优先级。
实现:
Advanteage:
Shortage:
Paper: MDPO: OVERCOMING THE TRAINING-INFERENCE DIVIDE OF MASKED DIFFUSION LANGUAGE MODELS
与上述方案不同,MDPO 并不改变位置选择的逻辑(依然沿用 Confidence-based 启发式),而是通过 RL 训练直接优化 Backbone 自身,使其输出的 Logit 能够更完美地适配推理时的调度轨迹。
实现
端到端优化: Backbone 参数 $\theta$ 是可训练的。在 RL 阶段,模型按照推理时的 Confidence-based 策略进行采样生成。
RL 模拟推理过程中的“递进式解掩码”路径。通过对比不同轨迹的累积奖励(如使用组相对优势估计 GRPO),强制模型在当前推理分布下产生更准确的预测。
RL 目标会拉大正确路径上 Token 的 Logit 值,并压低错误/干扰路径的值。这使得模型在推理时的 Confidence 评分不再仅仅是训练损失的副产品,而是真正具备了导航能力的启发式函数。
Advantage:
Shortage:
我们在这一部分介绍一些人为选定的scheduler
Paper: SPMDM: Enhancing Masked Diffusion Models through Simplifying Sampling Path
与传统的全局随机调度不同,SPMDM 认为高效的采样路径应具备两个特征:局部优先(优先解开已解 token 附近的掩码)和逻辑顺序优先(优先解开推理链条中靠前的子序列)。
Insight
Training Method:子序列划分与非均匀噪声
Target:通过子序列级的 NELBO 进行优化,强制模型在给定的上下文 $x_t^{-k}$ 下预测特定的子序列 $x^k$:
\[\mathcal{L}_{NELBO}=\sum_{k=1}^{K}\mathbb{E}_{t_{k}\sim[0,1],q_{t_{k}/0}}\left[\frac{\alpha_{t_{k}}^{\prime}}{1-\alpha_{t_{k}}}\log p_{\theta}(x^{k}\mid x_{t_{k}}^{k},{x_{t}}^{-k})\right]\]Paper: BLOCK DIFFUSION: INTERPOLATING BETWEEN AUTOREGRESSIVE AND DIFFUSION LANGUAGE MODELS
BD3-LMs 提出了一种在自回归(AR)建模与离散扩散建模之间进行插值的新范式。其核心动机在于通过引入块级(Block-wise)的结构,克服标准扩散模型在生成任意长度序列、推理效率(KV Caching)以及似然建模质量(Perplexity)上的局限性。
BD3-LM 将长度为 $L$ 的序列划分为 $B$ 个长度为 $L’$ 的块。模型在块之间遵循自回归分解,而在每个块内部执行离散扩散过程:
\[\log p_{\theta}(x) = \sum_{b=1}^{B} \log p_{\theta}(x^{b} \mid x^{\lt b})\]| 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,直接塑造了模型在训练与推理中所面对的条件分布,从而决定了三类核心误差的权衡关系:
分析表明,这三者并非独立问题,而是通过 Scheduler 的策略性选择在同一条误差曲线上被共同调节:
综合来看,RL-Policy Scheduler 通过reward,同时兼顾了内在序与全相关误差的制约,并且其采样路径始终是遵循policy的,更加训推一致,生成效果趋于最好。
Here are some more articles you might like to read next: