PaperlayerView exact revision on arXiv

Source-linked AI summary

Simplex Relaxation for Discrete Diffusion

Jinya Sakurai, Patrick Pynadath, Satoshi Hayakawa, Jaehong Yoon, Xulei Yang, Nancy F. Chen, Xun Xu

arXiv:2608.10615v1cs.CL

TL;DR

Uniform discrete diffusion leaves open whether training objectives and samplers can be enriched without changing the categorical corruption process. Simplax adds an exact Dirichlet–categorical augmentation, improving the OpenWebText perplexity–entropy tradeoff and achieving the strongest Sudoku results across evaluated clue densities and unconditional generation.

  • Problem

    Uniform discrete diffusion leaves open whether training objectives and samplers can be enriched while keeping the categorical corruption process unchanged.

  • Method

    Simplax augments uniform discrete diffusion with an auxiliary simplex state, yielding a Rao–Blackwellized reverse-bridge objective and stochastic sampler while preserving the categorical marginal.

  • Results

    Simplax achieves the strongest reported performance across OpenWebText generation and all evaluated Sudoku settings, including unconditional generation.

  • Takeaways & Limitations

    The results support enriching uniform categorical diffusion with auxiliary simplex structure for improved text perplexity–entropy tradeoffs and constrained-generation performance.

  • Takeaways & Limitations

    The formulation is specialized to uniform categorical corruption, with computational overhead not fully characterized and the concentration schedule remaining an additional design choice.

Abstract

from arXiv · show

Discrete diffusion models for categorical generation are defined by a corruption kernel, which determines the intermediate state space and the associated reverse prediction problem. We study uniform discrete diffusion and ask whether its training objective and reverse transitions can be enriched without changing the underlying categorical corruption process. We introduce Simplax, an exact Dirichlet--categorical augmentation that couples each corrupted categorical state with an auxiliary simplex-valued variable while preserving the original uniform diffusion process as its categorical marginal. This augmentation yields a tractable Rao--Blackwellized reverse-bridge objective and a corresponding stochastic reverse sampler, while retaining the corrupted categorical state as the denoiser input. Empirically, Simplax improves the generative perplexity--entropy tradeoff on unconditional OpenWebText generation. On Sudoku, a model trained exclusively on $30$-clue puzzles achieves the highest accuracy among the compared methods across all evaluated clue densities, including the minimum uniquely solvable $17$-clue regime, and also achieves the highest validity in unconditional generation.

1 Introduction

Simplax enriches uniform discrete diffusion with an auxiliary simplex-valued state while preserving the original categorical corruption process. It provides tractable reverse-bridge training and stochastic sampling, improving results on OpenWebText and Sudoku.

  • Method: Simplax introduces an exact Dirichlet–categorical augmentation that preserves uniform discrete diffusion as its categorical marginal.Each corrupted categorical state is coupled to an auxiliary simplex variable rather than replaced by one.
  • Method: The method derives a tractable categorical reverse-bridge surrogate with an analytically Rao–Blackwellized auxiliary expectation.The surrogate averages the standard discrete reverse KL over an auxiliary categorical decode from the simplex state.
  • Method: Simplax also yields a stochastic ancestral sampler derived from the same augmented hierarchy.The categorical state remains the denoiser input while the auxiliary simplex state supports reverse-transition construction.
  • Experiments: On OpenWebText, Simplax improves the generative perplexity–entropy tradeoff across a wide range of inference budgets.It outperforms the compared methods at most reported operating points.
  • Experiments: On Sudoku, Simplax achieves the highest performance among compared methods across all evaluated settings.Models are trained exclusively on 30-clue puzzles and evaluated in-distribution, across easier and harder clue densities, and in unconditional generation.

2 Preliminaries

The paper establishes one-hot categorical and simplex notation, defines the Dirichlet distribution, and specifies the standard discrete diffusion process and reverse-KL training objective. The noisy state mixes the clean token with a prior according to the noise schedule, while the model predicts the clean-token distribution.

  • Notation: Categorical variables are represented as one-hot vectors in V, while ΔK−1 denotes the probability simplex and Cat(·; π) the categorical distribution.Dirichlet-density results assume strictly positive class probabilities, satisfied by the uniform experimental base distribution.
  • Dirichlet distribution: Dir(·; ηp) is centered at p, with its normalized concentration vector giving the mean and η controlling concentration around that mean.The Dirichlet distribution is defined over the probability simplex.
  • Discrete diffusion: The noisy categorical state follows q(z_t | x) = Cat(z_t; p_t(x)), where p_t(x) := α_t x + (1 − α_t)π.The process uses prior π ∈ ΔK−1 and noise schedule α_t ∈ [0, 1].
  • Discrete diffusion: The schedule satisfies α_0 = 1, α_1 = 0, and α_t < 1 for t > 0, making p_t(x) strictly positive at positive times.These conditions support the positivity assumptions used for the noisy categorical distributions.
  • Reverse modeling: Training minimizes the categorical reverse KL D_KL[q(z_s | z_t, x) ∥ q(z_s | z_t, x̂_θ)], with x̂_θ = f_θ(z_t, t) predicting the clean-token distribution.The forward transition and corresponding reverse posterior have closed forms.

3 Method

Simplax augments uniform discrete diffusion with an exact Dirichlet–categorical hierarchy while preserving the categorical corruption process. The hierarchy enables a tractable Rao–Blackwellized reverse-bridge objective and a stochastic reverse sampler that retains the categorical state as denoiser input.

  • Hierarchy: Simplax adds an auxiliary simplex-valued variable without changing the categorical corruption process, forming an exact Dirichlet–categorical hierarchy.The construction derives reverse-bridge identities, a Rao–Blackwellized training objective, and a hierarchy-induced sampler.
  • Hierarchy: The relaxed state exactly decodes its categorical state by sampling zt from Cat(wt), while reverse posteriors remain analytically characterized.Given wt, x and zt are conditionally independent, and q(zt | wt, x) = q(zt | wt) = Cat(zt; wt).
  • Training objective: The direct simplex-bridge KL is generally intractable because its reverse conditional posterior is a Dirichlet mixture.Simplax instead averages the discrete reverse KL over an auxiliary decoder sample while retaining zt as the denoiser input.
  • Training objective: The auxiliary decoder expectation can be marginalized exactly, yielding a fully tractable Rao–Blackwellized objective without sampling noise from that decoder.The network prediction remains conditioned on the separately sampled categorical input zt.
  • Sampling: At inference, a stochastic ancestral sampler maintains (zt, wt), samples the next categorical state, and uses zt as the next network input.The auxiliary state wt carries bridge information needed for the next reverse posterior, while the final output is the categorical sample z0.

4 Experiments

Experiments evaluate Simplax on unconditional OpenWebText generation and cross-clue Sudoku generation. Simplax improves the OpenWebText perplexity–entropy frontier and achieves the strongest Sudoku performance across clue densities, including unconditional validity.

  • Design diagnostics: Self-conditioning has no budget-independent advantage: omitting it is better near data entropy at NFE = 16, whereas using it is better at NFE = 128.This diagnostic uses the w_t-input setting.
  • Design diagnostics: The z_t-input model attains lower Gen. PPL at comparable Gen. ENT at both evaluated NFE values, motivating its use in the main experiments.The main experiments retain w_t in the objective and reverse update while using z_t as the denoiser input.
  • Design diagnostics: UDLM initialization improves the Gen. PPL–Gen. ENT frontier at both NFE values compared with training Simplax from scratch.The comparison uses 800k UDLM iterations followed by 200k iterations with the Simplax objective, versus 1M Simplax iterations from scratch.
  • OpenWebText main comparisons: Simplax has the lowest Gen. PPL under all three evaluators at NFE = 16 and 1,024.At NFE = 128, Simplax is best under GPT-2 Large and GPT-2 XL, while LangFlow is best under Llama-2 7B.
  • Sudoku main comparisons: 95.85% validity is achieved by Simplax for unconditional Sudoku generation, versus 80.95% for the strongest baseline.All models are trained with 30 clues, and Simplax achieves the highest performance across conditional and unconditional settings.

5 Related Work

The paper situates Simplax within three research threads: categorical discrete diffusion, auxiliary-variable and hybrid formulations, and generative processes defined directly on the simplex.

  • Discrete diffusion for categorical data: Discrete categorical diffusion evolved from multinomial formulations to D3PM, which unified the framework and introduced structured corruption kernels and the standard variational training recipe.D3PM includes uniform and absorbing corruptions.
  • Auxiliary-variable and hybrid formulations: Auxiliary-variable and hybrid methods enrich discrete diffusion with correlated, latent, circuit-based, Gaussian-relaxed, or otherwise structured reverse formulations.Examples include Di4C, VADD, CoDD, Duo, and Duo+.
  • Diffusion and flow on the simplex: Simplex-based approaches instead define diffusion or flow directly on the simplex using continuous processes, categorical SDEs, Cox–Ingersoll–Ross dynamics, Dirichlet scores, or flow matching.The cited examples include simplex diffusion, DDSM, and Dirichlet Flow Matching.

6 Conclusion and Limitations

Simplax augments uniform discrete diffusion with an auxiliary simplex state while preserving its categorical forward process, improving reported OpenWebText and Sudoku results. The formulation is limited to uniform corruption, adds uncharacterized computational overhead, and leaves the concentration schedule as a design choice.

  • Contributions: Simplax is an exact Dirichlet–categorical augmentation that preserves uniform discrete diffusion’s categorical forward process while enabling a Rao–Blackwellized reverse-bridge objective and stochastic ancestral sampler.The auxiliary state is simplex-valued, and the categorical corrupted state remains the basis of the underlying process.
  • Empirical results: Simplax improves the Gen. PPL–Gen. ENT tradeoff on OpenWebText and achieves the highest Sudoku performance among compared methods across evaluated clue densities from 40 to 17 clues, including unconditional generation.The conclusion reports these results without providing numerical metric values.
  • Limitations: The formulation is specialized to uniform categorical corruption and introduces auxiliary simplex-state computational overhead whose magnitude relative to standard discrete diffusion remains uncharacterized.These limitations constrain the current method’s scope and efficiency assessment.
  • Limitations: The concentration schedule remains an additional design choice rather than being determined by the theory.The paper identifies this schedule dependence as a limitation of the present formulation.
  • Future directions: Future work should extend Simplax to broader categorical corruption kernels and develop more efficient reverse solvers.Both directions are stated as important next steps.

Supplementary Material · A Auxiliary Identities and Full Dirichlet–Categorical Hierarchy · A.1 Useful identities

The appendix establishes the exact probabilistic structure of the simplex relaxation at positive diffusion times. Its Dirichlet identities explain why categorical mixtures over one-hot shifts collapse to an unshifted Dirichlet density, making the relaxation exact rather than approximate.

  • A Auxiliary Identities and Full Dirichlet–Categorical Hierarchy: At positive diffusion times, p_t is strictly positive, and the hierarchy defines the clean endpoint as z_0 = x without introducing w_0.The Dirichlet hierarchy applies only at positive diffusion times.
  • A Auxiliary Identities and Full Dirichlet–Categorical Hierarchy: The auxiliary state w_t is introduced through a specified Dirichlet bridge rather than as a heuristic soft surrogate.This construction underpins the exact probabilistic treatment of the relaxation.
  • A Auxiliary Identities and Full Dirichlet–Categorical Hierarchy: The joint model has a closed hierarchy in both directions, with an exact Dirichlet marginal, exact decoding of z_t from w_t, and tractable reverse bridges after marginalizing z_t or z_s.These properties follow from the probabilistic structure developed in the appendix.
  • A.1 Useful identities: The first useful identity states that adding a one-hot count to a Dirichlet concentration vector multiplies the base density by the corresponding simplex coordinate.This is the basic shift formula for Dirichlet densities.
  • A.1 Useful identities: For concentrations α = ηp with p_k > 0 and η > 0, the specialized identity provides the cancellation used throughout the appendix.The corollary assumes p lies in the simplex interior and defines α = ηp.
  • A.1 Useful identities: A categorical mixture over one-hot shifts of a Dirichlet distribution collapses to the unshifted Dirichlet density, making the simplex relaxation exact rather than approximate.This collapse is the central consequence of Corollary 1.

A.2 Full Dirichlet–categorical hierarchy · B Rao–Blackwellized Reverse-Bridge Objective · C Continuous-Time Limit of the Main Objective

The appendix establishes a coherent Dirichlet–categorical hierarchy with exact decoding and closed-form relaxed reverse bridges. It then derives the Rao–Blackwellized objective used in training and continuous-time analysis.

  • A.2 Full Dirichlet–categorical hierarchy: The relaxed state has the correct Dirichlet marginal, admits exact categorical decoding, and preserves closed-form reverse bridges through simplex-conditioned posteriors.The reverse bridge becomes a mixture of shifted Dirichlet components after lifting the discrete reverse posterior.
  • A.2 Full Dirichlet–categorical hierarchy: Given wt, the discrete state zt is conditionally independent of x and can be recovered exactly as Cat(zt; wt).The decoder posterior is independent of x after applying the hierarchy’s marginal identities.
  • A.2 Full Dirichlet–categorical hierarchy: For s < t, the reverse posterior of zs given wt and x remains categorical, lifting the standard discrete reverse posterior to the relaxed state.This follows by combining the usual reverse posterior with the exact decoder q(zt | wt, x).
  • A.2 Full Dirichlet–categorical hierarchy: For 0 < s < t, the reverse posterior of ws given wt and x is a Dirichlet mixture with shifted components and lifted reverse-posterior weights.The shifted component form is Dir(ws; ηsps + ek), while the weights come from the lifted discrete posterior.
  • B Rao–Blackwellized Reverse-Bridge Objective: The Rao–Blackwellized objective analytically marginalizes the auxiliary decoder sample ezt while retaining the independently sampled denoiser input zt.The auxiliary sample satisfies q(ezt = ej | wt) = wt,j, and the resulting objective depends on wt through ⟨wt, log ˆpt⟩ and ρs|t(x, wt).
  • B Rao–Blackwellized Reverse-Bridge Objective: The closed-form objective is expressed through the current-time average ⟨wt, log ˆpt⟩ and the lifted reverse posterior ρs|t(x, wt).This is the expression used in the main objective and in the continuous-time analysis.
  • C Continuous-Time Limit of the Main Objective: The relaxed discrete bridge objective admits a non-degenerate first-order continuous-time limit, whose full proof is provided in the appendix.The limit is the first-order local limit stated in the main text.

C.1 Proof of the main continuous-time limit … D.3 Relations among the surrogate objectives

The appendix proves that the relaxed discrete bridge objective has the stated continuous-time limit, then organizes alternative tractable surrogates and derives their KL-based relations. The main-text objective is the categorical component of a richer joint bridge, while marginalization and averaging establish upper-bound relations among objectives.

  • C.1 Proof of the main continuous-time limit: C.1 establishes the first-order expansion of the relaxed bridge objective as Δℓct(wt, ˆxθ, x, t) + o(Δ).The derivation collects first-order terms, discards the θ-independent term −Δλ(t), and proves the continuous-time equations.
  • D Alternative Surrogate Objectives: D defines the exact relaxed bridge as a KL objective, but it is generally intractable because q(ws | wt, x) is a Dirichlet mixture.The tractable surrogates vary by whether they match zs or (zs, ws), and whether they condition directly on wt or decode ezt first.
  • D.1 A broader surrogate family: D.1 presents four tractable surrogates that separately choose categorical versus joint targets and direct versus decoded conditioning.The family includes Lzs|wt, the decoded categorical objective ¯Lzs|zt,wt, the direct joint objective Lzs,ws|wt, and the decoded joint objective ¯Lzs,ws|zt,wt.
  • D.1 A broader surrogate family: The loose joint objective decomposes into the main-text categorical objective plus an additional simplex-matching term, explaining why the selected objective has a non-degenerate continuous-time limit.The decomposition is ¯Lzs,ws|zt,wt = ¯Lzs|zt,wt + ¯Lws|zs,wt.
  • D.2 Auxiliary KL inequalities: D.2 supplies the KL tools used later: data processing, marginalization, and joint convexity of relative entropy.Marginalization cannot increase KL, and averaging related distributions is controlled by joint convexity.
  • D.3 Relations among the surrogate objectives: D.3 derives the relations Lzs|wt ≤ ¯Lzs|zt,wt ≤ ¯Lzs,ws|zt,wt and ¯Lws|zt,wt ≤ ¯Lzs,ws|zt,wt.These bounds follow from marginalizing zs or ws from joint bridges and averaging over the shared decoded-state distribution q(ezt | wt).
  • D.3 Relations among the surrogate objectives: The joint surrogate’s simplex-matching term has a closed form derived using Dirichlet expectation identities and comparison with KL between unshifted Dirichlet distributions.The appendix states Lemma 4 and Proposition 8 for this term, then derives the result by expanding Dirichlet densities and applying the identity.

D.4 Why the other surrogate objectives do not yield suitable continuous-time objectives

The tight discrete objective vanishes at second order in the local limit, while both joint objectives retain a generally nonzero O(1) simplex-matching term. This contrast motivates selecting the relaxed discrete bridge for the continuous-time formulation.

  • Local-limit comparison: Both joint objectives retain a generally nonzero O(1) simplex-matching term in the local limit.The tight joint and loose joint objectives are asymptotically given by the simplex term plus o(1).
  • Local-limit comparison: The tight discrete objective vanishes after division by ∆, so it does not yield a finite first-order training density.Its local expansion has a canceled first-order term, leaving second-order behavior.
  • Objective selection: The relaxed discrete bridge lies between these extremes and is selected as the natural objective for the continuous-time formulation.It avoids both the tight objective’s vanishing first-order scale and the joint objectives’ persistent O(1) simplex term.

E OpenWebText Experimental Details

The appendix specifies OpenWebText preprocessing, optimization, checkpoint, model-input, numerical-stability, and qualitative-generation protocols. It uses fixed 1,024-token blocks, detailed training controls, and entropy-matched examples evaluated at NFE = 16, 128, and 1,024.

  • Data preprocessing: OpenWebText documents are concatenated with end-of-sequence separators and packed into fixed-length blocks of 1,024 GPT-2 tokens.Both openwebtext-train and openwebtext-valid splits are used.
  • Optimization details: Training uses Adam with β1 = 0.9, β2 = 0.999, numerical constant 10^-8, gradient clipping at 1.0, bfloat16 precision, and a 2,500-step linear warmup.The learning rate is then constant, and generation uses parameters averaged with EMA decay 0.9999.
  • Simplax checkpoint: The Simplax checkpoint starts from a UDLM model trained for 800,000 steps and adds 200,000 Simplax-objective steps, totaling 1,000,000 optimization steps.This is the checkpoint used in the main OpenWebText comparison.
  • Model and objective: The network predicts clean-token distributions from categorical state z_t while retaining relaxed state w_t in training, with η = 0.01 and constant loss weighting.The reported checkpoint does not use auxiliary self-conditioning.
  • Numerical stability: Numerical safeguards use float64 for Dirichlet computations, restrict concentrations to [10^-10, 10^8], and softly bound logits with ℓ ← 30 tanh(ℓ/30).These restrictions are applied for numerical stability before normalization.
  • Qualitative generations: Representative generations at NFE = 16, 128, and 1,024 are selected using the main experiment’s entropy-matching procedure and are not manually rewritten.Excerpts are truncated at [...] and line-wrapped only for presentation; entropy is measured as generative unigram entropy in nats per token.

E.1 Sudoku experimental details

The Sudoku experiments train all models exclusively on 30-clue puzzles and evaluate conditional completion across clue densities from 40 to 17, plus an unconditional no-clue setting. Inputs use a fixed 180-token representation and shared Transformer configurations, while qualitative results show Simplax can satisfy Sudoku’s global constraints where baselines fail.

  • Dataset: The training set contains 48,000 uniquely solved puzzles with 30 observed cells, and every model is trained exclusively on this 30-clue dataset.Evaluation uses 2,000 puzzles for each clue count in the 25-or-more-clue settings.
  • Evaluation settings: Evaluation tests conditional completion with 40, 35, 30, 25, 20, and 17 clues, plus no-clue generation with all 81 prefix cells blanked.The 17-clue set uses uniquely solvable puzzles directly, while each 20-clue puzzle adds three solution clues.
  • Sequence representation: Each example contains 180 tokens: a BOS token, an 89-token puzzle, another BOS token, and an 89-token solution.The puzzle prefix remains 91 tokens at every clue count, with unobserved cells represented explicitly by a blank token.
  • Shared architecture: All methods use Transformer models with hidden dimension 512, eight blocks, eight attention heads of dimension 64, and dropout probability 0.1.The autoregressive model uses causal attention, while the remaining models use bidirectional attention with AdaLN-based time conditioning.
  • Qualitative Sudoku generation: In a representative 25-clue puzzle, Simplax satisfies the coupled row, column, and subgrid constraints, whereas baseline methods produce globally inconsistent grids.MDLM differs from the unique solution in eight cells, but Simplax produces a valid solution in the illustrated example.
Loading 2608.10615v1…