PaperlayerView exact revision on arXiv

Source-linked AI summary

Continuously Augmented Discrete Diffusion model for Categorical Generative Modeling

Huangjie Zheng, Shansan Gong, Ruixiang Zhang, Tianrong Chen, Jiatao Gu, Mingyuan Zhou, Navdeep Jaitly, Yizhe Zhang

arXiv:2510.01329v1stat.MLcs.LG

TL;DR

Masked discrete diffusion collapses unobserved possibilities into an information void, limiting semantic guidance during denoising. CADD pairs discrete masking with continuous latent diffusion so noisy embeddings guide token reconstruction, and reports consistent quality gains across text, image, and code generation while retaining simple, efficient training.

  • Problem

    Absorbing [MASK] states erase semantic information about masked tokens, limiting graded guidance for resolving ambiguity and maintaining coherence.

  • Method

    CADD augments discrete masking with paired continuous diffusion, using noisy latent vectors as soft semantic hints during reverse token denoising.

  • Results

    CADD consistently improves generative quality over mask-based diffusion across text, image, and code generation.

  • Takeaways & Limitations

    Continuous augmentation preserves semantic information while retaining discrete constraints, enabling more coherent generation and controllable diversity through latent estimation.

Abstract

from arXiv · show

Standard discrete diffusion models treat all unobserved states identically by mapping them to an absorbing [MASK] token. This creates an 'information void' where semantic information that could be inferred from unmasked tokens is lost between denoising steps. We introduce Continuously Augmented Discrete Diffusion (CADD), a framework that augments the discrete state space with a paired diffusion in a continuous latent space. This yields graded, gradually corrupted states in which masked tokens are represented by noisy yet informative latent vectors rather than collapsed 'information voids'. At each reverse step, CADD may leverage the continuous latent as a semantic hint to guide discrete denoising. The design is clean and compatible with existing discrete diffusion training. At sampling time, the strength and choice of estimator for the continuous latent vector enables a controlled trade-off between mode-coverage (generating diverse outputs) and mode-seeking (generating contextually precise outputs) behaviors. Empirically, we demonstrate CADD improves generative quality over mask-based diffusion across text generation, image synthesis, and code modeling, with consistent gains on both qualitative and quantitative metrics against strong discrete baselines.

1 Introduction

CADD addresses the information loss of masked diffusion by pairing discrete masking with continuous semantic diffusion. It preserves graded token information, supports diversity through latent resampling, and remains compatible with simple, parameter-efficient training.

  • 1 Introduction: CADD augments discrete masking with parallel continuous diffusion, letting noisy latent vectors provide semantic hints while discrete context constrains denoising.The paired representation replaces collapsed masked states with graded information and combines strengths of continuous and masked diffusion.
  • 1 Introduction: Masked states retain graded proximity to the ground-truth embedding, reducing ambiguity and making discrete token prediction easier.The latent representation decays smoothly rather than losing information abruptly.
  • 1 Introduction: Multi-sample latent estimation can explore alternative valid token or span choices, enhancing generation diversity.Resampling continuous latents at inference provides multiple plausible semantic hints for a discrete state.
  • 1 Introduction: CADD retains standard token cross-entropy and continuous diffusion losses, while sampling can alternate or jointly update discrete and continuous states.This keeps the training and sampling design compatible with existing discrete diffusion procedures.
  • 1 Introduction: CADD uses no special architecture, matches prior masked-diffusion parameter counts, and adds no significant training compute cost.The same backbone can be fine-tuned from existing masked-diffusion checkpoints.

2 Related Work

Related work includes discrete diffusion with masking and continuous relaxations for categorical data. These approaches motivate CADD’s attempt to preserve semantic structure without losing categorical constraints or balancing mode seeking and mode covering.

  • Discrete Diffusion Models: Discrete diffusion models define Markov chains over token space using uniform or absorbing transitions, with later work enriching masking and reverse-process operations.Related techniques address masking representation, remasking, and editing within discrete diffusion.
  • Continuous Relaxations for Discrete Data: Continuous approaches denoise token embeddings or relax tokens into Euclidean or simplex spaces, but unconstrained relaxations can fail to preserve categorical discreteness.These methods provide continuous representations while raising constraints around categorical structure.
  • Bridging Through the Lens of Mode Balancing: Prior work balances mode seeking and mode covering through guidance, score distillation, and related diversity–precision control techniques.This line of work frames generation as a trade-off between sharper samples and broader coverage.

3 Preliminary

The preliminary formulation represents categorical sequences in discrete token space and continuous embedding space. It defines discrete corruption and reverse prediction alongside Gaussian continuous diffusion and its denoising objective.

  • Discrete Diffusion Models: A categorical sequence uses a vocabulary containing V tokens plus a dedicated mask token m, with each token position represented in the discrete state space.The notation distinguishes ordinary vocabulary symbols from the absorbing mask symbol.
  • Continuous Diffusion Models: Tokens are mapped from one-hot discrete vectors into learnable d-dimensional embeddings, producing continuous representations z_0 = w_θ(x_0).The embedding matrix w_θ maps each discrete token to a continuous vector and the sequence to z_0 ∈ R^{n×d}.
  • Discrete Diffusion Models: The discrete forward process independently corrupts tokens according to a time-dependent transition schedule, while the reverse model predicts original data from corrupted states.The schedule α_t is strictly decreasing, and reverse prediction is trained through a variational negative-log-likelihood bound.
  • Continuous Diffusion Models: Continuous diffusion corrupts real-valued embeddings with Gaussian noise and trains a reverse process using an SNR-reweighted mean-squared-error objective.The forward process has a closed-form Gaussian distribution, while the reverse network estimates denoising transitions.

4 Continuously Augmented Discrete Diffusion (CADD)

CADD augments masked discrete diffusion with paired continuous latent diffusion, preserving graded semantic information and using it to guide reverse token prediction. Its training combines discrete prediction with continuous modeling, while sampling can use multiple latent estimates to represent alternative token choices.

  • 4 Continuously Augmented Discrete Diffusion (CADD): CADD pairs the discrete token process with a continuous latent variable that preserves semantic hints when tokens are masked.The joint construction factorizes discrete and continuous components, with the discrete path triggering the latent evolution.
  • 4.2 Reverse: CADD predicts token distributions by averaging predictions conditioned on multiple sampled continuous latents, allowing plausible continuous states to inform discrete denoising.This expectation over latent states is intended to represent possible tokens more accurately and can support diverse predictions.
  • 4.1 Forward: When a token is first masked, its embedding begins Gaussian degradation and then becomes progressively noisier while remaining masked, instead of losing its signal abruptly.The discrete path acts as a trigger: unmasked embeddings remain unchanged, whereas masked embeddings follow a smooth degradation trajectory.
  • 4.3 Algorithm and Implementation: CADD implementation adds noisy masked-position embeddings to discrete embeddings before the backbone produces per-position categorical logits.Training forms clean embeddings, applies continuous noise only at masked positions, fuses the two representations by element-wise addition, and optimizes cross-entropy.
  • 4.3 Algorithm and Implementation: The model can be trained with standard cross-entropy, while an additional reweighted MSE term may more closely estimate the exact variational objective.The simplified cross-entropy objective is used in most experiments because it is more computationally efficient.
  • 4.3 Algorithm and Implementation: During sampling, the chain starts with fully masked tokens and continuous Gaussian latents, then repeatedly predicts tokens or updates latents across reverse steps.The sampling procedure iterates from timestep T to 1 and maintains the coupled discrete-continuous state.

5 Experiments

CADD is evaluated on text, image, and code generation against discrete diffusion baselines, showing consistent gains across modalities and sampling settings.

  • 5 Experiments: CADD consistently improves over strong discrete diffusion baselines across text, image, and code generation experiments.The evaluations cover OpenWebText, CIFAR-10, ImageNet-32, and code-generation benchmarks.
  • 5 Experiments: At T = 4096, CADD’s MAUVE increases by 0.3 from T = 1024, while generative perplexity decreases from 44.6 to 35.3.Masked-only baselines stagnate or degrade over this sampling range.
  • 5 Experiments: CADD uses 168M trainable parameters and matches MDLM’s 0.5h inference time for 5000 samples on 8 H100 GPUs when K = 1.The design adds only forward and fusion computation at K = 1; cost increases linearly with K.
  • 5 Experiments: CADD reaches CIFAR-10 FID 2.88 and Inception Score 10.04 with 512 NFE, surpassing MDM variants by 0.38 FID.On ImageNet-32, CADD obtains FID 3.74 and outperforms all reported baselines.
  • 5 Experiments: CADD improves DiffuCoder on HumanEval from 67.1 to 72.0 and BigCodeBench-Hard from 12.8 to 17.6.Fine-tuning from a DiffuCoder checkpoint further improves HumanEval to 73.8 and BigCodeBench to 41.5.

6 Conclusion

CADD augments absorbing discrete diffusion with an auxiliary continuous space that preserves semantic information after masking and guides reverse denoising.

  • 6 Conclusion: CADD retains semantic information in a smooth continuous representation after discrete absorption, enabling more coherent and contextually accurate generations.The paper reports higher-quality samples across text, image, and code generation.

A.1 ELBO Derivation

The appendix derives the ELBO and reverse-process structure for coupled discrete and continuous diffusion, using tractable factorization and case-specific continuous posteriors.

  • A.1 ELBO Derivation: The ELBO decomposes into prior matching, reverse-step terms, and a decoder term, becoming tight when the reverse conditionals match the true posteriors and the prior match is zero.The final decoder must equal the true conditional induced by the joint distribution.
  • A.1 ELBO Derivation: The joint marginal q(x_t, z_t | x_0) factorizes into q(x_t | x_0) and q(z_t | x_t, x_0), with both components tractable in closed form.This factorization supports the subsequent ELBO and posterior derivations.
  • A.1 ELBO Derivation: For remaining masked positions, the continuous reverse posterior is Gaussian with parameters obtained by multiplying two Gaussian factors.Unmasked positions yield a Dirac posterior, while the remaining-masked case follows the standard DDPM Gaussian derivation.
  • A.1 ELBO Derivation: The adopted posterior factorization is tractable because the alternative ordering produces a complex Gaussian mixture and lacks a closed-form discrete conditional.The implementation therefore uses the factorization from Proposition 2 for efficient sampling.
  • A.1 ELBO Derivation: The forward chain couples discrete masking with continuous diffusion, using separate schedules for the two components.The discrete process follows α(t) = 1 − t, while masked continuous latents follow a linear flow-matching path to isotropic noise.

B.2 Experiment-Specific Settings

Experiments cover text, image, and code generation with modality-specific datasets, model settings, and sampling evaluations. CADD is also tested as a fine-tuning objective for an existing masked diffusion checkpoint.

  • Experiment setup: CADD is evaluated across text, image, and code generation using OpenWebText, CIFAR-10/ImageNet, and OpenCoder settings.Text uses a 168M-parameter DiT trained on OpenWebText; image experiments use 32×32 pixel-space tokens; code experiments adapt a 7B Qwen2.5-Coder checkpoint.
  • Experiment setup: Text generation uses 5,000 samples of length 1,024 and evaluates MAUVE, generative perplexity, and entropy across 128 to 4,096 sampling steps.
  • Fine-tuning: Figure 5 compares the fine-tuned CADD-objective checkpoint with CADD and the initialized MDLM checkpoint.
  • Fine-tuning: CADD is additionally evaluated by fine-tuning an existing MDLM checkpoint with the CADD objective for 50B additional tokens.

C.2 Perplexity Evaluation

Perplexity evaluation reports CADD’s performance on standard text benchmarks and zero-shot datasets while noting that its mixed discrete-continuous objective limits direct bound-tightness comparisons. CADD is competitive with strong diffusion baselines and achieves the best diffusion results on several evaluations.

  • Evaluation caveat: Direct comparison of bound tightness is considered unfair because CADD’s objective includes KL divergence terms for both discrete and continuous components.
  • Benchmark perplexity: On LM1B, CADD achieves the best results among diffusion models for discrete-part perplexity on both LM1B and OpenWebText data.
  • Benchmark perplexity: On Text8, CADD’s perplexity is highly competitive but slightly worse than the state-of-the-art RDLM.
  • Zero-shot evaluation: CADD achieves the best diffusion-model zero-shot perplexity on 4 of 7 tasks, while CADD and MDLM surpass autoregressive models on Lambada, Pubmed, and Arxiv.

C.3 Ablation studies

Ablations show that more continuous samples and more sampling steps improve CADD quality, while estimator choice controls mode-seeking versus mode-covering behavior. Architecture and fusion choices also affect evaluation outcomes, with ADM performing best among tested image architectures.

  • Sampling samples and K: Increasing sampling steps and K consistently improves CADD performance, with the largest generative-perplexity gain occurring from K=2 to K=3.Larger K improves x̂0 estimation but increases inference compute K times, creating a quality-latency trade-off.
  • Generation variance: CADD has the lowest entropy among the highest-quality text-generation models, while larger K further lowers entropy without eliminating strong MAUVE performance.
  • Fusion choice: Fusion choices produce only 0.03 absolute MAUVE and 0.07 absolute entropy variation, although concatenation requires an additional projection layer.
  • Estimator choice: Hard x̂0 estimation raises MAUVE by 0.06 and lowers entropy by 0.11, whereas soft estimation yields higher entropy and more diversity.
  • Image architecture: ADM outperforms DDPM++ across tested NFEs, with ADM at NFE=512 showing especially strong image-generation performance and the best qualitative visual quality.

D.1 Text Samples

The passages present qualitative CIFAR-10 results alongside descriptions of configurable input-output transitions and an unrelated serotonin comparison study.

  • CIFAR-10 qualitative results are shown for ADM at NFE values of 64, 256, and 512.The figure presents these settings from the top row to the bottom row.
  • A reported randomized comparison examined serotonin-related effects alongside aerobic stimuli and sleep stimulation.The passage states that measurable differences were observed, but its description is incomplete.
  • The extension supports custom functions and configurable transitions across inputs and outputs.The passage describes different input-output settings and transition controls.
  • Another passage describes changes in performance, balance, healing activity, and concentration fluctuations under varying conditions.The relationships among these conditions are not clearly specified in the passage.

D.2 Code Samples

The section presents code samples generated on HumanEval, MBPP, and BigcodeBench, including sorting, rolling-maximum, and histogram-generation programs.

  • The rolling-maximum sample scans numbers while tracking the current maximum and appends each updated maximum to the result.
  • Generated samples include HumanEval, MBPP, and BigcodeBench code examples.
  • The sorting sample uses a gap-based loop and swaps elements when an earlier value exceeds the later comparison value.
  • The MBPP sample appends 12, derives a random-list size from the input sum, generates integers from 1 to 100, and plots their histogram.
  • The MBPP function validates that my_list is a list of numeric elements and raises TypeError or ValueError otherwise.
Loading 2510.01329v1…