Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Seunghwan Jang, SooJean Han

TL;DR

Stratified Hazard Sampling (SHS) introduces a minimal-variance event-scheduling rule for CTMC/DTMC discrete diffusion and flow models, replacing independent per-step jumps with hazard-space stratification driven by a single random phase. SHS preserves the expected number of edits and the model’s multimodality while bounding jump-count variance to at most $1/4$, addressing under-edit and over-edit failure modes common with uniform-noise initializations. A phase-allocation extension biases early edits for high-risk lexical constraints, reducing late-masking artifacts without retraining. Across a few-step MNIST reconstruction task and a lightweight UDLM text-generation protocol, SHS improves stability and sample quality in low-NFE regimes, indicating practical benefits for fast, reliable discrete-sample generation in real-world applications.

Abstract

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.

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

TL;DR

Stratified Hazard Sampling (SHS) introduces a minimal-variance event-scheduling rule for CTMC/DTMC discrete diffusion and flow models, replacing independent per-step jumps with hazard-space stratification driven by a single random phase. SHS preserves the expected number of edits and the model’s multimodality while bounding jump-count variance to at most , addressing under-edit and over-edit failure modes common with uniform-noise initializations. A phase-allocation extension biases early edits for high-risk lexical constraints, reducing late-masking artifacts without retraining. Across a few-step MNIST reconstruction task and a lightweight UDLM text-generation protocol, SHS improves stability and sample quality in low-NFE regimes, indicating practical benefits for fast, reliable discrete-sample generation in real-world applications.

Abstract

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.
Paper Structure (65 sections, 2 theorems, 30 equations, 4 figures, 1 table, 2 algorithms)

This paper contains 65 sections, 2 theorems, 30 equations, 4 figures, 1 table, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. CTMC/DTMC-based Discrete Generative Models
  4. CTMC formulation.
  5. DTMC formulation (discrete diffusion).
  6. Unified jump-mass view.
  7. Lexical Safety as Conditioning, and Why Naive Masking Is Biased
  8. Exact conditioning via Doob $h$-transform.
  9. Naive masking creates error w.r.t. $p_1(x\mid\mathcal{A})$.
  10. Related Works
  11. Discrete Diffusion and CTMC Generative Models
  12. Iterative Masked-Token Refinement for Non-Autoregressive Generation
  13. Constrained and Safe Generation via Lexical Constraints
  14. Simulation of Inhomogeneous Poisson Processes and Markov Jump Processes
  15. Stratified Hazard Sampling
  16. ...and 50 more sections

Key Result

Proposition 4.1

Fix $S\ge 0$ and draw $\theta\sim\mathrm{Uniform}(0,1)$. Let $\{S\}=S-\lfloor S\rfloor$ denote the fractional part and define $J=\lfloor S\rfloor+\mathbf{1}[\theta<\{S\}].$ Then $\mathbb{E}[J]=S$.

Figures (4)

  • Figure 1: Standard step-based sampler vs. SHS (ours).
  • Figure 2: Hazard-space jump locations. Histograms of $H_k$ (the cumulative hazard at the $k$-th jump). EM sampling exhibits Erlang/Gamma-shaped variability, while SHS produces stratified, bounded-support locations, demonstrating reduced trajectory randomness.
  • Figure 3: Quantized MNIST reconstructions under varying step budgets. EM vs. SHS at $n \in \{256,64,16,8\}$ discretization steps. SHS reduces the under-edit regime and yields cleaner early-stage reconstructions.
  • Figure 4: Gen. PPL vs. NFE (UDLM protocol). SHS improves generation quality most notably in the few-step regime, indicating reduced sampler-induced instability under tight compute budgets.

Theorems & Definitions (2)

  • Proposition 4.1: Unbiasedness
  • Proposition 4.2: Minimal variance