PaperlayerView exact revision on arXiv

Source-linked AI summary

Simple Guidance Mechanisms for Discrete Diffusion Models

Yair Schiff, Subham Sekhar Sahoo, Hao Phung, Guanghan Wang, Sam Boshar, Hugo Dalla-torre, Bernardo P. de Almeida, Alexander Rush, Thomas Pierrot, Volodymyr Kuleshov

arXiv:2412.10193v3cs.LG

TL;DR

Controllable guidance for discrete diffusion is difficult because continuous methods do not directly apply, and some discrete models cannot repeatedly edit tokens. The paper adapts classifier-based and classifier-free guidance and develops uniform-noise diffusion with a tightened continuous-time variational bound. It reports improved controllable generation across several discrete domains and strong uniform-noise language-model performance on small-vocabulary data.

  • Problem

    Continuous guidance methods do not directly apply to discrete diffusion, and some discrete diffusion models cannot repeatedly edit generated tokens.

    Source: §1 · p. 1

  • Method

    The paper adapts classifier-based and classifier-free guidance for discrete diffusion and develops uniform-noise models with a tightened continuous-time variational lower bound.

    Source: §1 · pp. 1–2, §4.2 · p. 6

  • Results

    Discrete guidance improves controllable generation relative to autoregressive and prior diffusion guidance methods across several domains, while uniform-noise diffusion reaches state-of-the-art performance on small-vocabulary datasets and among uniform-noise models.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    Uniform-noise diffusion can be amenable to repeated editing and guidance while achieving strong performance on small-vocabulary data.

    Source: §1 · p. 2, §4 · p. 5, §6 · p. 10

    Extending classifier-based guidance to sequences is not straightforward and requires an assumption about classifier conditioning on the sequence.

    Source: §3.2 · p. 4

Abstract

from arXiv · show

Diffusion models for continuous data gained widespread adoption owing to their high quality generation and control mechanisms. However, controllable diffusion on discrete data faces challenges given that continuous guidance methods do not directly apply to discrete diffusion. Here, we provide a straightforward derivation of classifier-free and classifier-based guidance for discrete diffusion, as well as a new class of diffusion models that leverage uniform noise and that are more guidable because they can continuously edit their outputs. We improve the quality of these models with a novel continuous-time variational lower bound that yields state-of-the-art performance, especially in settings involving guidance or fast generation. Empirically, we demonstrate that our guidance mechanisms combined with uniform noise diffusion improve controllable generation relative to autoregressive and diffusion baselines on several discrete data domains, including genomic sequences, small molecule design, and discretized image generation.

1 INTRODUCTION

The introduction identifies obstacles to controllable discrete diffusion and presents guidance methods plus uniform-noise diffusion as a response. Experiments report improved controllable generation across several domains and strong language-modeling performance for uniform-noise diffusion on small-vocabulary data.

1 INTRODUCTION

  • Discrete data poses three challenges: continuous guidance gradients are undefined, some diffusion extensions cannot revise tokens repeatedly, and controllable uniform-noise models lag autoregressive models in perplexity.The paper notes that repeated editing can help correct errors during generation.Source: §1 · p. 1
  • The paper adapts classifier-based and classifier-free guidance to discrete diffusion and revisits uniform-noise models that can repeatedly edit samples.It derives a continuous-time evidence lower bound to tighten the variational gap and improve uniform-noise model performance.Source: §1 · pp. 1–2
  • Across genomics, molecule strings, and discretized images, discrete diffusion is reported as more controllable than autoregressive models with classifier-free guidance; classifier-based guidance improves on prior diffusion and autoregressive guidance methods.The paper also reports state-of-the-art performance on small-vocabulary datasets and a new perplexity state of the art among uniform-noise diffusion models.Source: §1 · p. 2

2 BACKGROUND

The background establishes notation and describes discrete diffusion as corruption followed by learned denoising. It then introduces classifier-based and classifier-free guidance as ways to condition diffusion sampling.

2 BACKGROUND

  • The paper represents tokens as one-hot vectors over a vocabulary, sequences as lists of tokens, and categorical distributions by their probability vectors.A diffusion model adds noise to latent variables through a fixed corruption process and learns to remove it.Source: §2 · p. 2, §2.1 · p. 2

2.1 DISCRETE DIFFUSION MODELS

  • D3PM defines discrete corruption through transition matrices, while a simpler class interpolates between clean data and a noisy prior using a decreasing noise schedule.For absorbing-state diffusion, unmasked tokens remain unmasked at earlier times; the paper contrasts this with other noise choices later.Source: §2.1 · pp. 2–3

2.2 DIFFUSION GUIDANCE

  • Classifier-based guidance samples from a tempered distribution combining a pretrained diffusion model with a classifier score, weighted by an inverse-temperature parameter.The classifier term and diffusion-model term are combined in the guided distribution.Source: §2.2 · p. 3
  • Classifier-free guidance replaces the classifier score with the difference between conditional and unconditional diffusion-model scores, then combines those scores using the guidance weight.The conditional and unconditional models are trained separately as conditioned and unconditioned diffusion models.Source: §2.2 · p. 3

3 GUIDANCE ALGORITHMS FOR DISCRETE DIFFUSION

The section derives classifier-free and classifier-based guidance for discrete diffusion, where guidance terms are not differentiable with respect to discrete representations. The methods use tempered reverse distributions and token-level sequence strategies to make guidance tractable.

3.1 CLASSIFIER-FREE GUIDANCE

  • Classifier-free guidance combines conditional and unconditional reverse probabilities with weights γ and 1 − γ in log space.The paper names this method D-CFG.Source: §3.1 · pp. 3–4
  • For sequences, D-CFG models conditional and unconditional denoisers together by randomly dropping or masking the conditioner during training.The resulting distributions factorize independently across sequence positions, enabling efficient sampling from the tempered guidance distribution.Source: §3.1 · p. 4

3.2 CLASSIFIER-BASED GUIDANCE

  • Classifier-based guidance tempers the reverse transition using the classifier probability, and its single-token distribution can be normalized over N possible token values.The method is named D-CBG.Source: §3.2 · pp. 4–5
  • Sequence-level classifier guidance assumes classifier outputs factorize across tokens, avoiding normalization over all N^L possible sequences.Without factorization, the required normalization could involve N^L terms.Source: §3.2 · p. 4
  • Restricting classifier-based guidance to sequences differing at only one token makes normalization require a sum over N terms.The method adapts FUDGE guidance for autoregressive generation to discrete diffusion.Source: §3.2 · pp. 4–5
  • Direct sequence classifier guidance can require O(L · N) classifier passes per denoising step, while a first-order Taylor approximation uses one forward and backward pass.The approximation treats the classifier as a continuous function of one-hot inputs.Source: §3.2 · p. 5

4 UNIFORM DIFFUSION LANGUAGE MODELS

Uniform-noise diffusion permits repeated edits to intermediate tokens, which the paper uses to develop models more amenable to guidance. Its continuous-time likelihood-bound derivation tightens the objective, and sequence models trained with the refined objective are called UDLMs.

4 UNIFORM DIFFUSION LANGUAGE MODELS

  • Uniform noise lets an intermediate token change more than once during denoising, unlike masked diffusion where an unmasked token remains unmasked.The paper motivates this repeated refinement as making uniform-noise diffusion more amenable to guidance.Source: §4 · p. 5, §4.1 · p. 5

4.1 UNIFORM NOISE DIFFUSION

  • The denoising model predicts clean data from a noisy latent and uses that prediction to parameterize the reverse process.The predictive model is xθ(zt, t), with output in the N-category probability simplex.Source: §4.1 · p. 5

4.2 IMPROVED LIKELIHOOD BOUNDS IN CONTINUOUS TIME

  • Taking T → ∞ yields a simpler closed-form NELBO, reduces reconstruction and prior losses to zero, and further tightens the diffusion loss through a continuous-time extension.The derivation analyzes the reconstruction, diffusion, and prior terms separately.Source: §4.2 · p. 6
  • When α_t(T) = 0 and pθ(z_t(T)) = π, the prior-loss KL divergence evaluates to zero.The corruption process then has terminal distribution q(z_t(T)) = π.Source: §4.2 · p. 6
  • Sequence training assumes denoising factorizes independently across tokens given noisy latents, and models trained with the refined sequence objective are called UDLMs.The passage identifies the objective as sequence-level.Source: §4.2 · p. 6

5 EXPERIMENTS

The experiments evaluate language modeling across discrete domains and test guided generation in genomics, molecule design, and images. They find strong uniform-noise language modeling in smaller-vocabulary settings and improved controllability with discrete diffusion guidance.

5.1 LANGUAGE MODELING WITH UNIFORM NOISE DISCRETE DIFFUSION

  • UDLM attains state-of-the-art performance on small-vocabulary datasets, with performance near MDLM on text8 but a more persistent gap at larger vocabularies.Table 1 reports perplexity across datasets; the authors relate the larger-vocabulary gap to the larger set of possible clean sequences uniform-noise models must predict.Source: §5.1 · p. 7, Table 1 · p. 7
  • On text8 and LM1B, UDLM achieves the best reported language-modeling performance among uniform-noise discrete diffusion models.Tables 2 and 3 report the text8 and LM1B comparisons, respectively.Source: §5.1 · p. 7, Table 2 · p. 8, Table 3 · p. 8

5.2 GUIDED DISCRETE DIFFUSION

  • For genomic sequences, diffusion models achieve higher species-control F1 scores than AR, and UDLM improves control as guidance strength rises while AR and MDLM metrics degrade.Diffusion generation uses T = 512, compared with AR decoding 32,768 tokens.Source: §5.2 · p. 8, Table 4 · p. 8
  • Across guided tasks, classifier-free guidance is more useful with diffusion than with autoregressive models, while D-CBG is the strongest discrete classifier-based method, especially with UDLM.Figure 3 highlights D-CBG versus FUDGE for drug-likeness and diffusion D-CFG versus AR for the novelty–ring-count trade-off.Source: §5.2 · p. 7, Figure 3 · p. 9
  • For QM9 molecule generation, MDLM and UDLM suit ring-count maximization better under D-CFG, while D-CBG better balances property maximization with valid, novel outputs than AR or NOS.Table 5 evaluates validity, novelty, and mean QED or ring count for novel generated sequences.Source: §5.2 · p. 9, Table 5 · p. 9
  • On CIFAR10, MDLM and UDLM outperform finite-time D3PM on FID and Inception Score, with further improvement when D-CFG is added.Table 6 compares finite-time and continuous-time models.Source: §5.2 · pp. 9–10, Table 6 · p. 9
  • With faster CIFAR10 sampling at smaller T, UDLM is robust while MDLM performance deteriorates; UDLM can revise tokens throughout decoding.The authors contrast this with MDLM predictions that can become locked in.Source: §5.2 · p. 10

6 RELATED WORKS, DISCUSSION, AND CONCLUSION

This section discusses UDLM’s discrete diffusion guidance and continuous-time variational bound, including a derivation step toward its ELBO equivalence with SEDD.

6 RELATED WORKS, DISCUSSION, AND CONCLUSION

  • UDLM extends variational discrete diffusion to uniform noise, while its guidance is derived directly for discrete data rather than through continuous representations.The authors distinguish their approach from absorbing-state methods and from concurrent guidance work tailored to continuous-time Markov chain rate matrices, including Nisonoff et al. (2024).Source: §6 · p. 10
  • The continuous-time bound closes the performance gap to absorbing-state diffusion, while uniform noise is on par or better in small-vocabulary regimes.The authors also report that classifier-free guidance enables control without sacrificing sequence quality, and their classifier-based method outperforms previous methods for autoregressive and diffusion models.Source: §6 · p. 10

A CONTINUOUS TIME DISCRETE UNIFORM DIFFUSION

  • Taking T →∞ yields a continuous-time diffusion-loss formulation for uniform noise, assuming the process interpolates between clean data and noise.The limiting distribution is uniform, u = 1/N, where N is the vocabulary size.Source: A CONTINUOUS TIME DISCRETE UNIFORM… · p. 17

B RELATING GUIDANCE TO CTMC

  • The appendix relates the paper’s guidance mechanisms to continuous-time Markov chains and states that their continuous-time extensions are equivalent to formulas derived using that framework.The comparison is with concurrent work by Nisonoff et al. (2024).Source: B RELATING GUIDANCE TO CTMC · p. 19

B.1 CTMC FORMULATION OF D-CFG

  • Discrete classifier-free guidance is proportional to sampling with the predictor-free guided reverse rate from Nisonoff et al. (2024), using forward Euler discretization.The derivation uses the paper’s discrete guidance formulation and the CTMC reverse-rate formulation.Source: §B.1 · pp. 20–21

B.2 CTMC FORMULATION OF D-CBG

  • Discrete classifier-based guidance is equivalent to predictor guidance in Nisonoff et al. (2024); under process factorization, classifier evaluation is needed only for candidates differing in at most one token.The continuous-time extension is proportional to sampling with the corresponding reverse rate matrix and Euler discretization.Source: §B.2 · pp. 21–22

C RELATING UDLM TO CTMC

C.1 EQUIVALENCE OF UDLM ELBO AND SEDD ELBO

  • For a specific denoising-score-matching parameterization, the SEDD and UDLM negative evidence lower bounds are equivalent.The derivation uses UDLM’s mean parameterization, which predicts clean data from noisy observations.Source: §C.1 · pp. 23, 25
  • The derivation proceeds by using equations (53) and (54).The excerpt does not include the resulting expression or further details.Source: §C.1 · p. 25

D.1 DATASET DETAILS

The appendix documents dataset preparation and model architectures across genomic, molecular, image, and language experiments. It specifies dataset splits, tokenization, and the model backbones used.

D.1 DATASET DETAILS

  • Species10 uses ten reference genomes chunked into 32,768-nucleotide segments; the data is base-pair tokenized and split randomly 95%/5%.Species labels condition guidance; Table 10 summarizes species proportions and split sizes.Source: §D.1 · p. 25, Table 10 · p. 27
  • QM9 contains about 133k molecules represented as canonical SMILES, with QED and ring-count labels and a 95%/5% random split.Sequences use 32-token inputs; guidance labels use the 90th-percentile cutoff for QED or ring count.Source: §D.1 · p. 26
  • CIFAR10 uses 50,000 training and 10,000 validation images, tokenized as integer pixel intensities from 0 to 255.Guidance is conditioned on image class labels.Source: §D.1 · p. 26
  • text8 uses lowercase letters and whitespace in non-overlapping 256-token chunks, with the first 90M characters for training and final 5M for validation.Source: §D.1 · p. 26
  • Amazon Review and LM1B use bert-base-uncased tokenization and 128-token padded sequences; their training/validation sizes are 3.6M/400k and 7M/72k, respectively.LM1B sequences are chunked before the split sizes are reported.Source: §D.1 · pp. 26–27

D.2 ARCHITECTURAL DETAILS

  • Genomics uses Mamba-based backbones, CIFAR10 uses a UNet, and QM9 and NLP use a 12-layer Diffusion Transformer.The genomic diffusion model uses Caduceus-Ph with 4.8M parameters versus 3.5M for the AR model; the CIFAR10 UNet adds a discretized truncated logistic transformation and label embedding.Source: §D.2 · pp. 27–28

D.3 TRAINING CONFIGURATIONS

This appendix gives training settings, describes guidance baselines and implementations, and defines the genomic sequence evaluation metrics.

D.3 TRAINING CONFIGURATIONS

  • All diffusion models were trained and evaluated with a log-linear noise schedule; for CIFAR10, our models trained for 300K iterations, compared with 1.5M for D3PM (Austin et al., 2021).Source: §D.3 · p. 28

D.4.1 BASELINES

  • FUDGE perturbs autoregressive next-token probabilities using a classifier, while PPLM updates hidden representations with classifier gradients and a KL penalty.FUDGE truncates candidate tokens to the top-k by unconditional likelihood; PPLM combines perturbed and unperturbed distributions.Source: §D.4.1 · pp. 28–29
  • NOS guidance uses a classifier built from a frozen pretrained diffusion backbone, with only its final linear projection trained.Inference follows Langevin sampling with a stability/fluency KL penalty.Source: §D.4.1 · p. 29

D.4.2 D-CFG DETAILS

  • D-CFG training randomly masks the class condition at a 10% rate and inference uses conditional and unconditional model passes.Source: §D.4.2 · p. 29

D.4.3 D-CBG DETAILS

  • D-CBG classifiers are trained on noised inputs using a log-linear schedule whose corruption type matches the diffusion model being guided.For QM9, the classifier uses a smaller Diffusion Transformer with mean pooling and a linear projection to classification logits.Source: §D.4.3 · p. 30

D.5.1 GENOMIC SEQUENCES METRICS

  • Genomic evaluation uses weighted 3-mer and 6-mer distribution-shift scores, discriminator AUROC for real-versus-generated sequences, and oracle macro-F1 for species control.The discriminator trains on generated and validation sequences; the separate oracle classifier is trained on the Species10 dataset.Source: §D.5.1 · p. 30

D.5.2 CIFAR10 QUALITY METRICS

The appendix describes quality and control metrics for CIFAR10 image generation and QM9 molecule guidance.

D.5.2 CIFAR10 QUALITY METRICS

  • CIFAR10 image quality is evaluated with FID and Inception Score, using 50,000 generated images per model.FID compares real and generated feature distributions; Inception Score reflects class-specificity and diversity.Source: §D.5.2 · pp. 30–31
  • A pretrained CIFAR10 classifier’s F1 score serves as a proxy for whether generated images satisfy the requested condition.Source: §D.5.2 · p. 31

D.5.3 QM9 GUIDANCE METRICS

  • QM9 guidance evaluation reports SMILES validity, novelty among valid unique molecules, and mean and median values for the targeted QED or ring-count property.The experiments generate 2,048 sequences of length 32.Source: §D.5.3 · p. 31

E.1 EFFECT OF VARYING GUIDANCE HYPERPARAMETERS

Guidance hyperparameter sweeps show that outcomes depend on model and task: UDLM is more amenable to genomic guidance, while increasing γ has mixed effects on CIFAR10. Removing D-CBG’s first-order approximation improves QED guidance for both diffusion models, but ring-count results differ between UDLM and MDLM.

E.1 EFFECT OF VARYING GUIDANCE HYPERPARAMETERS

  • UDLM achieves better quality and control metrics than AR and MDLM for species-specific genome generation, and its performance does not degrade as γ increases.AR and MDLM performance degrades as γ increases; the comparison varies γ over {1, 2, 3} and diffusion sampling steps over {128, 256, 512}.Source: §E.1 · p. 32, Table 15 · p. 32
  • Increasing γ generally improves IS and F1 for both MDLM and UDLM on CIFAR10, while FID improves for MDLM but worsens for UDLM when γ exceeds 2.The sweep uses γ values 1–5 with T = 1000; visual samples become more refined and sharpened as γ increases.Source: §E.1 · pp. 32–33

E.2 EFFECT OF USING FIRST ORDER APPROXIMATION IN D-CBG

  • Removing D-CBG’s first-order approximation boosts QED maximization performance for both MDLM and UDLM.The comparison is presented for both models in Figure 5.Source: §E.2 · p. 33
  • For ring-count maximization, removing the approximation improves UDLM’s novelty and ring counts, while MDLM becomes unstable for γ > 1 and can decode the full sequence early without recovering from mistakes.The first-order approximation for MDLM does not appear to show the same instability.Source: §E.2 · p. 33

E.3 EFFECT OF VARYING T

Varying the number of decoding steps affects sample quality differently across domains and models. More steps are associated with better genomic quality for MDLM and UDLM, whereas UDLM is reported to handle faster CIFAR10 inference better than MDLM.

E.3 EFFECT OF VARYING T

  • Increasing T is positively related to genomic sample quality for both MDLM and UDLM, even though the tested T values are orders of magnitude below sequence length L = 32768.The tested values are T ∈ {128, 256, 512}.Source: §E.3 · p. 33
  • On QM9, MDLM benefits significantly from increasing T for both QED and ring-count guidance, while UDLM results are more consistent across decoding-step counts.For D-CBG, the analysis omits the first-order approximation for QED and uses it for ring count.Source: §E.3 · pp. 33–34
  • For CIFAR10, Table 35 reports additional T values beyond Table 7 in Section 5.2, and UDLM accommodates the faster inference setting better than MDLM owing to its ability to recover from ‘mistakes.’Source: §E.3 · p. 34

F UNCONDITIONAL SAMPLE GENERATION

The section evaluates unconditional samples from models trained on LM1B. It reports generative perplexity and provides random generated examples.

F UNCONDITIONAL SAMPLE GENERATION

Loading 2412.10193v3…