PaperlayerView exact revision on arXiv

Source-linked AI summary

Score Mismatch and Set-Constrained Sampling in Zero-Shot Conditional Discrete Diffusion

Renxiang Huang, Yuchen Liang, Yingbin Liang, Ness Shroff, Lifeng Lai

arXiv:2610.21061v1stat.MLcs.LG

TL;DR

Score mismatch complicates zero-shot conditional sampling with an unconditional discrete diffusion model, motivating analysis of its added error and methods for set constraints. The paper derives KL guarantees and proposes Randomized Hard Clamping, whose assignment distribution can be reconstructed from unconditional scores under absorbing rates.

Abstract

from arXiv · show

Discrete diffusion models have become a powerful framework for modeling complex discrete distributions across a range of applications. While theoretical convergence guarantees depend on consistency between the target and training distributions, such a case rarely holds in practice, making score mismatch a central challenge in discrete diffusion sampling. This is particularly relevant to zero-shot conditional generation, where a model trained on the unconditional distribution is directly reused at inference time to sample from a conditional target distribution. In this paper, we provide the first theoretical analysis for score-mismatched discrete diffusion models and derive a general KL divergence upper bound that clearly identifies the additional bias caused by such a mismatch. Our analysis covers partial-observation constraints, including text infilling, sentence completion, and more general coordinatewise set constraints. Specifically, we analyze the classical Conditional SEDD sampler and show that it can incur a bias term that scales with the dimension, vocabulary size, and compatibility. Based on this analysis, we propose a Randomized Hard Clamping (RHC) sampler for coordinatewise set constraints. We show that the assignment probabilities can be expressed exactly by the unconditional scores with absorbing rate. Finally, numerical and real-data experiments demonstrate the effectiveness of the proposed RHC sampler.

I. INTRODUCTION

The introduction frames zero-shot conditional discrete diffusion as sampling with mismatched scores, then presents a KL-based analysis and a randomized sampler for coordinatewise set constraints.

I. INTRODUCTION

  • Zero-shot conditional generation reuses an unconditional model at inference, but its score can mismatch the conditional target; this setting includes sequence completion and text infilling.The paper also considers coordinatewise set constraints, where each constrained position admits several candidate values.Source: §I · pp. 1–2
  • The paper asks how to characterize score-mismatch bias and extend fixed-observation sampling to set constraints while controlling approximation error.Under set constraints, an assignment is a feasible combination of values, and selecting one allows reuse of a fixed-observation sampler.Source: §I · p. 2

A. Our Contributions

  • The analysis gives a non-asymptotic KL upper bound separating intrinsic sampling error from additional bias caused by using training-distribution scores for a mismatched target.For discrete diffusion, mismatch perturbs the reverse CTMC rate matrix; the analysis also uses a score-entropy Bregman divergence rather than the continuous setting's usual L2 error.Source: A. Our Contributions · p. 3
  • For Conditional SEDD, hard-clamp and true conditional scores coincide on unobserved coordinates under absorbing rates, but uniform rates can introduce bias dependent on dimension, vocabulary size, and compatibility.Under absorbing rates, the convergence bound has no additional score-mismatch bias term; under uniform rates, the paper bounds that bias assuming condition compatibility.Source: A. Our Contributions · p. 3
  • Randomized Hard Clamping represents set-constrained targets as mixtures over assignments and reconstructs the exact assignment distribution from unconditional scores for absorbing-rate CTMCs.The paper reports that numerical simulations and real-data experiments support its theoretical results and the sampler's practical effectiveness.Source: A. Our Contributions · p. 3

B. Related Works

  • Prior zero-shot discrete diffusion work emphasized algorithm design and empirical performance, while this paper establishes non-asymptotic guarantees for absorbing and uniform rates and extends the framework to set constraints.The paper positions its general framework as quantifying additional sampling error from score mismatch.Source: B. Related Works · p. 4

A. (Score-matched) Discrete Diffusion Models

  • The setup models data on a discrete state space, gradually perturbs samples with a forward CTMC, and learns a reverse process for denoising.It considers absorbing rates that send tokens to MASK and uniform rates that distribute tokens uniformly at the terminal time.Source: A. (Score-matched) Discrete Diffusion… · pp. 4–5

B. Score Mismatch in Discrete Diffusion Models

The paper defines discrete score mismatch as a distortion between target and training reverse dynamics, then applies this framework to zero-shot conditional sampling. In that setting, an unconditional model is reused to sample unobserved coordinates given fixed observations, without additional training.

B. Score Mismatch in Discrete Diffusion Models

  • Discrete score mismatch modifies the target reverse rate relative to the training reverse rate by a multiplicative factor, 1 + Δ_t(y, x).The paper defines mismatch through a ratio of target and training scores, rather than as an additive perturbation.Source: B. Score Mismatch in Discrete… · p. 6
  • Unlike additive mismatch in continuous diffusion, discrete score mismatch distorts the reverse jump rate through changes in distribution ratios.This motivates defining mismatch as a ratio in the discrete setting.Source: B. Score Mismatch in Discrete… · p. 7

C. Zero-Shot Conditional Sampling under Fixed Observations

  • Zero-shot conditional sampling reuses a pretrained unconditional diffusion model to sample from the conditional target without additional training for the observations.For fixed observations, the target is the conditional distribution of unobserved coordinates given the observed values.Source: C. Zero-Shot Conditional Sampling… · p. 7

D. Zero-Shot Conditional Sampling under Set Constraints

The section extends zero-shot sampling from fixed observations to coordinatewise feasible sets and develops convergence guarantees for score-mismatched discrete diffusion. Its bounds separate oracle sampling error from mismatch-induced bias, with the bias depending on the CTMC rate choice.

D. Zero-Shot Conditional Sampling under Set Constraints

  • Coordinatewise set constraints allow each constrained coordinate to take values from a feasible set, and the target is the marginal conditioned on those constraints.Singleton feasible sets recover fixed observations; sampling uses a pretrained unconditional model without additional training.Source: D. Zero-Shot Conditional Sampling… · p. 7

III. GUARANTEES FOR GENERAL SCORE-MISMATCHED DISCRETE DIFFUSION MODELS

  • The general analysis assumes controlled score-estimation error along the target process and bounded score estimates.The stated assumptions are described as commonly used; the dependence on the score-estimate bound is logarithmic.Source: §III · p. 8

A. Absorbing-Rate CTMCs with τ-leaping

  • For absorbing-rate CTMCs with τ-leaping, the KL bound separates oracle sampling error from an additional bias that accumulates from score mismatch and vanishes when mismatch is zero.The oracle term includes initialization, discretization, and score-estimation errors; finite KL uses a surrogate initialization distribution [16, Theorem 1].Source: A. Absorbing-Rate CTMCs with τ-leaping · pp. 8–9

B. Extension to Uniform-Rate CTMCs

  • For uniform-rate CTMCs, the final KL bound likewise decomposes into oracle error and additional score-mismatch bias.The oracle error can use applicable score-matched analyses, including [17, Theorem 2].Source: B. Extension to Uniform-Rate CTMCs · p. 9
  • Under uniform-rate CTMCs, the score-mismatch bias bound is worse by O(S) than under absorbing-rate CTMCs because transitions can reach any of the S − 1 other states.The uniform-rate score-function bound therefore depends on vocabulary size.Source: B. Extension to Uniform-Rate CTMCs · p. 9

C. Extension to Uniformization Sampling

The paper extends its mismatch analysis to uniformization sampling and applies the framework to Conditional SEDD under fixed observations. Uniformization replaces the discrete-time accumulation of estimation error with a continuous-time integral, while hard clamping can satisfy observations yet bias unobserved coordinates.

C. Extension to Uniformization Sampling

  • Uniformization analysis requires score-estimation error to be controlled throughout continuous time, unlike τ-leaping analysis, which controls error on the selected time grid.The distinction follows the continuous-time versus time-discretized sampling methods considered.Source: C. Extension to Uniformization Sampling · p. 10
  • With absorbing rates, uniformization exactly simulates the sampling process without time-discretization error, while score mismatch accumulates as a continuous-time integral.The accumulated mismatch affects the KL divergence between target and generated distributions.Source: C. Extension to Uniformization Sampling · p. 10

IV. GUARANTEES FOR ZERO-SHOT CONDITIONAL DISCRETE DIFFUSION SAMPLERS UNDER FIXED OBSERVATIONS

  • Conditional SEDD hard-clamps observed coordinates at each reverse step, but this intermediate-state constraint may bias sampling on unobserved coordinates.The hard clamp enforces the observation on the diffusion state, whereas the true condition applies to clean data.Source: §IV · pp. 10–11

A. Conditional SEDD under Fixed Observations with Absorbing Rate

Conditional SEDD’s hard-clamp bias depends on the diffusion rate: absorbing-rate sampling adds no score-mismatch bias on unobserved coordinates, while uniform-rate sampling incurs a compatibility- and dimension-dependent term. For set constraints, RHC selects an assignment and then generates the remaining coordinates using hard clamping; absorbing-rate unconditional scores recover the target assignment probabilities exactly.

A. Conditional SEDD under Fixed Observations with Absorbing Rate

  • With absorbing-rate CTMC, Conditional SEDD adds no score-mismatch bias on unobserved coordinates; its KL bound contains only initialization, score-estimation, and time-discretization errors.The hard-clamp score coincides with the true conditional score under the stated assumptions.Source: A. Conditional SEDD under Fixed… · pp. 11–12

B. Conditional SEDD under Fixed Observations with Uniform Rate

  • The uniform-rate compatibility assumption requires a positive lower bound on context likelihood across completions, excluding completions to which the context assigns zero probability.For text generation, some completions may be preferred, but others must remain possible.Source: B. Conditional SEDD under Fixed… · p. 12
  • Under uniform-rate CTMC, Conditional SEDD accumulates score-mismatch bias over the unobserved coordinates, unlike its absorbing-rate result.The bound’s additional bias grows linearly with |Ω| for fixed γ; it does not decrease as the number of reverse steps increases.Source: B. Conditional SEDD under Fixed… · pp. 12–13
  • As compatibility γ approaches 1, the uniform-rate bias disappears; as γ approaches 0, the error upper bound diverges.γ quantifies compatibility between the fixed context and possible completions.Source: B. Conditional SEDD under Fixed… · pp. 12–13

V. RANDOMIZED HARD CLAMPING UNDER SET CONSTRAINTS

A. Assignment distribution and RHC Sampler

  • With absorbing-rate CTMC, unconditional scores recover the exact target assignment distribution for any diffusion time t > 0.This provides the assignment probabilities without requiring the true data distribution as an input.Source: A. Assignment distribution and RHC… · pp. 13–14
  • RHC samples an assignment from w(a), then hard-clamps its constrained coordinates while generating the remainder with pretrained unconditional scores.The learned-score distribution w approximates the target assignment distribution.Source: A. Assignment distribution and RHC… · p. 14

B. Convergence Guarantee for RHC

RHC’s convergence analysis separates assignment-distribution error from the sampling error for each selected assignment. With exact unconditional scores, the assignment error vanishes; otherwise its additional bound scales with the square root of the number of constrained coordinates and has no explicit dependence on the feasible-set size.

B. Convergence Guarantee for RHC

  • RHC’s convergence guarantee controls both assignment-distribution error and sampling error for each selected assignment.The latter reuses the fixed-observation convergence analysis for each assignment.Source: B. Convergence Guarantee for RHC · pp. 14–15
  • The assignment-distribution error is controlled through score-estimation accuracy along the paths used to construct w.The assumption averages estimation error under both the target distribution v and estimated distribution w; with exact scores, the error vanishes.Source: B. Convergence Guarantee for RHC · p. 15
  • With exact unconditional scores, w = v and the assignment-distribution error vanishes; otherwise, its additional bound scales with √|Ω̄| and has no explicit dependence on |A|.The number of feasible sequences may grow combinatorially with the number of constrained values.Source: B. Convergence Guarantee for RHC · p. 15

VI. EXPERIMENT RESULTS

The experiments test Conditional SEDD under fixed observations and RHC under set constraints. RHC lowers perplexity across six text scenarios with a small generation-time increase; Conditional SEDD’s behavior depends on the rate and constraint compatibility.

A. Conditional SEDD under Fixed Observation

  • With absorbing rate, Conditional SEDD closely tracks the oracle as both KL errors decrease with more reverse steps, without additional structural score-mismatch bias.The hard-clamped score coincides with the true conditional score on the observed coordinates.Source: Figure 1 · p. 16, A. Conditional SEDD under Fixed… · p. 16
  • With uniform rate, Conditional SEDD reaches an irreducible KL plateau while the oracle achieves lower KL, and finer time discretization does not eliminate their mismatch.Both samplers’ KL divergences gradually decrease as reverse steps increase.Source: A. Conditional SEDD under Fixed… · p. 16, Figure 2 · p. 17
  • Under uniform-rate Conditional SEDD, KL generally increases with more unobserved coordinates in three tested observations, while another observation has the worst KL.The bound depends on both the number of free coordinates and compatibility; changing the observation from [0, 2, 2] to [0, 2, 1] lowers compatibility and increases KL.Source: A. Conditional SEDD under Fixed… · pp. 16–17, Figure 3 · p. 17

B. RHC under Set Constraints

  • For RHC under absorbing rate with exact assignment probabilities and unconditional scores, KL decreases with more steps and remains comparable to the oracle.In this synthetic experiment, the assignment distribution has no approximation error, and hard clamping introduces no score mismatch for each assignment.Source: B. RHC under Set Constraints · pp. 17–18, Figure 4 · p. 18
  • RHC achieves lower perplexity than the uniform-assignment baseline across six constraint scenarios, with total generation time increasing by 3.1%–3.4%.Table I reports means ± standard deviations and total generation time across five independent trials.Source: Table I · p. 18, B. RHC under Set Constraints · p. 18

VII. CONCLUSION

  • The conclusion extends the analysis to coordinatewise set constraints through RHC with absorbing rate and identifies more general multimodal data and complex constraints as future work.The paper reports hard-clamp score equality with the true conditional score under absorbing rate, but not under uniform rate.Source: §VII · p. 19

Appendix

The appendix situates zero-shot diffusion among prior approaches and develops a KL-error decomposition for score-mismatched sampling. It separates oracle sampling error from the additional contribution of score mismatch.

APPENDIX I RELATED WORKS

  • Prior zero-shot methods reuse unconditional diffusion without extra training, while prior theoretical score-mismatch guarantees cited here concern continuous states.The appendix also distinguishes specialized conditional frameworks from approaches that control inference using unconditional models.Source: Appendix I · p. 23

APPENDIX II PROOF OF THEOREM III.3

  • The proof decomposes KL error into oracle error and a separate score-mismatch contribution, treating the target substitution as a multiplicative correction to the rate matrix.The oracle error comprises initialization, discretization, and score-estimation errors.Source: APPENDIX II PROOF OF THEOREM III.3 · pp. 23–24
  • For the absorbing-rate bound, the proof uses a score bound of (T − t_k)^−1 and early stopping to control the accumulated terms.The cited lemma bounds both score functions between zero and (T − t_k)^−1 under its stated positivity conditions.Source: APPENDIX II PROOF OF THEOREM III.3 · pp. 25, 27

APPENDIX IV PROOF OF THEOREM III.5

These proofs adapt convergence bounds to uniform-rate sampling and establish conditional-score relationships for fixed observations. They distinguish the absorbing-rate equality from the uniform-rate mismatch.

APPENDIX IV PROOF OF THEOREM III.5

  • For the uniform-rate convergence proof, the authors retain the proof structure of Theorem III.3 and replace the score-function and rate-matrix bounds with uniform-rate counterparts.The bias term is derived after taking expectations and summing over time steps.Source: APPENDIX IV PROOF OF THEOREM III.5 · pp. 27–28

APPENDIX V PROOF OF THEOREM III.7

  • Under absorbing rate, the hard-clamp score equals the true conditional score on unobserved coordinates, supporting the fixed-observation convergence bound.The appendix states this equality for indices k in Ω.Source: Appendix V · p. 30

APPENDIX VII PROOF OF THEOREM IV.3

  • For uniform rate, the proof represents hard-clamp and true conditional scores using distinct distributions, then bounds their mismatch.The appendix compares the distributions before applying the uniform-rate convergence result.Source: APPENDIX VII PROOF OF THEOREM IV.3 · p. 31

APPENDIX VIII

The proof separates set-constraint sampling error into assignment-distribution error and fixed-observation sampling error. The assignment term measures deviation from the target distribution over feasible assignments.

APPENDIX VIII

  • The proof defines the RHC distribution through a distribution w over feasible assignments, followed by hard-clamped sampling conditional on the selected assignment.The hard-clamped component refers to sampling with the fixed observation set to that assignment.Source: APPENDIX VIII · p. 33
  • Theorem VIII.1 separates total KL divergence under set constraints into assignment-distribution error KL(v∥w) and fixed-observation sampling error.The mixture decomposition does not require an absorbing rate.Source: APPENDIX VIII · p. 34
  • For each assignment, the fixed-observation convergence result controls the second error term in the decomposition.The proof uses a common oracle bound for fixed-observation sampling across assignments.Source: APPENDIX VIII · pp. 34–35

APPENDIX IX DIFFERENCE BETWEEN HARD CLAMP SCORE AND TRUE CONDITIONAL SCORE

A binary example compares the hard-clamp score with the true conditional score under a fixed observation. Under uniform rate, the two scores do not coincide.

APPENDIX X DETAILED SETTINGS OF EXPERIMENTS

The appendix specifies the synthetic and text experiment settings, discusses uniform-rate RHC reconstruction, and includes proof material for the absorbing-rate assignment result.

APPENDIX X DETAILED SETTINGS OF EXPERIMENTS

  • The synthetic experiments use a four-coordinate, four-token Gibbs distribution with token preferences and neighboring-symbol rewards.Absorbing-rate simulations add a MASK state.Source: A. Conditional SEDD under Fixed… · p. 37

A. Conditional SEDD under Fixed Observations

  • In Figures 1 and 2, coordinates 1–2 are observed at y = [0, 2]; Conditional SEDD is compared with an oracle conditional-score baseline based on Monte-Carlo estimates.For Figure 3's comparison, changing [0, 2, 2] to [0, 2, 1] lowers γ from approximately 0.02 to 0.004.Source: A. Conditional SEDD under Fixed… · p. 37

B. RHC under Set Constraints

  • For real-data text experiments, fluency and rationality are evaluated by full-sequence perplexity using pretrained GPT-2-large; lower perplexity indicates closer alignment with natural language and better generation quality.Text scenarios use sequence length 50, six constrained token positions, scenario-specific fixed prefixes, and a period suffix.Source: B. RHC under Set Constraints · p. 38

APPENDIX XI DISCUSSION OF RHC WITH UNIFORM RATE

  • Uniform-rate RHC can reconstruct the exact assignment distribution, but direct implementation presents computational and numerical challenges.The direct computation requires summing over [S]^d and inverting a matrix over the constrained space; absorbing rate avoids those additional steps.Source: APPENDIX XI DISCUSSION OF RHC WITH… · pp. 38–39

APPENDIX XII PROOF OF THEOREM V.1

APPENDIX XIV PROOF OF SUPPORTING LEMMAS

The supporting-lemma appendix gives proof steps for the fixed-observation and set-constraint convergence results, including a decomposition over feasible assignments.

A. Proof of Theorem VII.1

  • The proof of Theorem VII.1 distinguishes cases using the observed-set size and the number U of unmatched tokens between y and the observed coordinates.Source: A. Proof of Theorem VII.1 · p. 42

B. Proof of Theorem VIII.1

Related papers

Loading 2610.21061v1…