Stratified Hazard Sampling: Minimal-Variance Event Scheduling for CTMC/DTMC Discrete Diffusion and Flow Models

📄 arXiv: 2601.02799v1 📥 PDF

作者: Seunghwan Jang, SooJean Han

分类: cs.LG, cs.CL

发布日期: 2026-01-06

备注: Work in progress. Feedback welcome

💡 一句话要点

提出分层风险抽样(SHS)，最小化CTMC/DTMC离散扩散模型的事件调度方差，提升生成质量。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 离散扩散模型 CTMC DTMC 分层抽样 序列生成 方差最小化 事件调度

📋 核心要点

  1. 现有基于CTMC/DTMC的离散生成模型在均匀噪声初始化下，由于token独立决策导致编辑方差大，易出现欠编辑或过度编辑问题。
  2. 论文提出分层风险抽样(SHS)，通过分层累积风险来确定token的编辑事件，从而最小化编辑方差，提高生成质量。
  3. SHS无需超参数调整，且能保留多模态性。同时，针对词汇约束，提出了相位分配变体，缓解了后期掩蔽问题。

📝 摘要（中文）

本文提出了一种名为分层风险抽样(Stratified Hazard Sampling, SHS)的推理原则，用于基于CTMC/DTMC的离散生成模型，包括均匀噪声离散扩散（如D3PM/CTDD）和离散流匹配。这些模型通过时间非齐次马尔可夫过程重复替换token来实现非自回归序列生成。通常的推理方法是基于步进模拟：每个token在每个离散化步骤中通过独立的伯努利（或分类）抽样决定是否跳跃。在均匀噪声初始化下，由于自校正需要多次编辑每个位置，这些独立决策导致编辑次数和时间上的巨大方差，从而导致欠编辑（残留噪声）或过度编辑（级联不必要的替换）等典型失效模式，降低了可重复性。SHS将每个token的编辑建模为由累积风险（CTMC）或累积跳跃质量（DTMC）驱动的事件，并通过分层该累积量来放置事件：对于每个位置的单个随机相位，当其累积风险超过单位间隔阈值时，token就会跳跃。这保留了预期的跳跃次数，同时实现了无偏整数估计器中可能的最小方差（以1/4为界），而不会改变每次跳跃的目的地抽样，从而保留了多模态性。我们还介绍了一种用于黑名单式词汇约束的相位分配变体，该变体优先在高风险位置进行早期编辑，以减轻后期掩蔽伪影。

🔬 方法详解

问题定义：论文旨在解决基于CTMC/DTMC的离散生成模型在推理过程中，由于token独立决策导致的编辑次数和时间上的高方差问题。这种高方差会导致欠编辑（残留噪声）或过度编辑（级联不必要的替换）等问题，最终影响生成质量和可重复性。现有方法缺乏对编辑事件的有效控制，导致模型性能不稳定。

核心思路：论文的核心思路是将每个token的编辑过程视为由累积风险（CTMC）或累积跳跃质量（DTMC）驱动的事件。通过对累积风险进行分层抽样，确保在期望的跳跃次数不变的情况下，最小化编辑事件的方差。这种方法的核心在于利用分层抽样来控制编辑事件的发生，从而提高生成过程的稳定性。

技术框架：SHS的整体框架可以概括为以下几个步骤：1. 对于每个token，计算其累积风险或累积跳跃质量。2. 为每个位置分配一个随机相位。3. 根据累积风险和随机相位，确定token的跳跃事件。具体来说，当累积风险超过单位间隔阈值时，token发生跳跃。4. 根据模型原有的跳跃目的地抽样方法，确定跳跃的目标token。5. 对于词汇约束，使用相位分配变体，优先在高风险位置进行早期编辑。

关键创新：SHS最重要的创新点在于其最小化编辑事件方差的能力。通过分层抽样，SHS能够保证在期望跳跃次数不变的情况下，实现无偏整数估计器中可能的最小方差（以1/4为界）。与现有方法相比，SHS无需额外的超参数调整，并且能够保留模型原有的多模态性。此外，针对词汇约束的相位分配变体也是一个重要的创新，能够有效缓解后期掩蔽问题。

关键设计：SHS的关键设计在于累积风险的分层抽样。具体来说，对于每个token，首先计算其累积风险函数H(t)。然后，为每个位置分配一个均匀分布的随机相位U。token在时间t发生跳跃的条件是H(t) >= k + U，其中k为整数。这种设计保证了期望的跳跃次数不变，并且最小化了跳跃事件的方差。对于词汇约束，论文提出了一种相位分配策略，优先在高风险位置分配较小的相位值，从而使得这些位置更容易发生早期编辑。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文提出的SHS方法在理论上证明了其最小化编辑事件方差的能力，并通过实验验证了其在实际应用中的有效性。虽然摘要中没有给出具体的性能数据，但强调了SHS能够解决欠编辑和过度编辑问题，提高生成质量和可重复性。相位分配变体能够有效缓解词汇约束下的后期掩蔽问题。

🎯 应用场景

该研究成果可广泛应用于序列生成任务，例如文本生成、语音合成、蛋白质序列设计等。通过降低生成过程中的方差，可以提高生成质量和可控性，尤其是在需要精确控制编辑次数和位置的场景下，例如文本编辑和代码生成。

📄 摘要（原文）

CTMC/DTMC-based discrete generative models, including uniform-noise discrete diffusion (e.g., D3PM/CTDD) and discrete flow matching, enable non-autoregressive sequence generation by repeatedly replacing tokens through a time-inhomogeneous Markov process. Inference is typically implemented with step-based simulation: each token decides to jump via independent Bernoulli (or categorical) draws at every discretization step. Under uniform-noise initialization, where self-correction requires multiple edits per position, these independent decisions induce substantial variance in both the number and timing of edits, leading to characteristic failure modes such as under-editing (residual noise) or over-editing (cascading unnecessary substitutions), decreasing reproducibility. We propose Stratified Hazard Sampling (SHS), a drop-in and hyperparameter-free inference principle for any sampler that admits a stay-vs.-replace decomposition. SHS models per-token edits as events driven by cumulative hazard (CTMC) or cumulative jump mass (DTMC) and places events by stratifying this cumulative quantity: with a single random phase per position, a token jumps whenever its accumulated hazard crosses unit-spaced thresholds. This preserves the expected number of jumps while achieving the minimum possible variance among unbiased integer estimators (bounded by 1/4), without altering per-jump destination sampling and thus retaining multimodality. We also introduce a phase-allocation variant for blacklist-style lexical constraints that prioritizes early edits at high-risk positions to mitigate late-masking artifacts.