Dirichlet Diffusion Score Model for Biological Sequence Generation

1. Introduction

• Score-based generative SDEs support likelihood evaluation through a probability-flow ODE and flexible sampling, but the original formulation is not directly suitable for discrete data.The introduction situates this gap alongside prior discrete-diffusion approaches, noting that none used the continuous-time SDE framework except quantization-based methods.Source: §1 · p. 1
• The proposed continuous-time diffusion in the probability simplex uses a Dirichlet stationary distribution to model discrete data.Forward diffusion starts at simplex vertices and remains in the continuous simplex, bridging discrete data and continuous diffusion through a latent-variable interpretation.Source: §1 · p. 1
• DDSM generates Sudoku puzzles, solves hard puzzles without additional training, and designs diverse human promoter sequences comparable to genome promoters.The paper also reports competitive performance on binarized MNIST.Source: §1 · p. 1

2.1. Score-based Generative Modeling with SDE

• Reverse-time SDEs and the probability-flow ODE can be sampled using the score of the time-dependent data distribution, so learning reverse diffusion becomes score-function learning.The probability-flow ODE gives the same distribution at time t as the reverse-time SDE.Source: §2.1 · p. 2

2.2. Univariate Jacobi Diffusion Process

• Jacobi diffusion evolves within [0,1] and has a Beta(a,b) stationary distribution determined by its positive parameters a and b.Its diffusion coefficient approaches zero at either boundary while the drift points inward; its transition density has a spectral expansion.Source: §2.2 · p. 2

3.1. Forward Diffusion SDE for Two-Category Data

• For two-category data, Jacobi diffusion moves values from 0 or 1 continuously through [0,1] toward a Beta stationary distribution, enabling score-based forward and reverse diffusion.The model is trained with denoising score matching using the transition density.Source: §3.1 · pp. 2–3

3.2. Forward Diffusion SDE for k-Category Data

• For k categories, stick-breaking combines k−1 independent Jacobi diffusions into a simplex-valued process with Dirichlet stationarity, defining the Dirichlet diffusion score model.The construction is a finite-dimensional version of the GEM process; the authors note alternative Dirichlet-stationary processes are more computationally expensive as forward diffusions.Source: §3.2 · pp. 3–4

3.3. Score Matching Training for k-Category Discrete Data

• Training starts from one-hot data, transforms it to stick-breaking coordinates, samples Jacobi transitions or stationary Beta variables, and computes score-matching loss in v space.The learned model takes all v dimensions as input after conversion to x space for reverse sampling.Source: §3.3 · p. 4

3.4. Weighting Function for Score Matching Loss of General SDE

• The proposed weighted score-matching loss is invariant under bijective differentiable changes of variables, unlike the unweighted loss.When the diffusion coefficient is scalar or diagonal and independent of v, the formulation recovers likelihood weighting.Source: §3.4 · p. 4

3.5. Likelihood Computation

• The probability-flow ODE provides continuous-space likelihood computation, which can be converted to likelihood in x space by a change-of-variable formula.The derivation uses the stationary distribution and an unbiased Hutchinson estimate of the Jacobian trace.Source: §3.5 · p. 4, LIKELIHOOD COMPUTATION FOR CONTINUOUS… · p. 5

BOUNDING DISCRETE DATA LIKELIHOOD WITH VARIATIONAL LOWERBOUND

• A discrete-data likelihood lower bound marginalizes over simplex-valued categorical parameters and estimates its ELBO by sampling forward diffusion near time zero.The bound tightens as the starting time approaches zero; the training objective also reduces the variational gap, which approaches a KL divergence up to a constant.Source: LIKELIHOOD COMPUTATION FOR CONTINUOUS… · p. 5, BOUNDING DISCRETE DATA LIKELIHOOD WITH… · p. 5

3.6. Improving Sampling Efficiency and Sample Quality

• An efficiency strategy reduces forward sampling for k categories from k−1 Jacobi processes to effectively one univariate Jacobi process.The strategy is described in the appendix.Source: §3.6 · p. 5
• Time-dilation adjusts reverse sampling toward high-probability regions without changing the score model, for flat stationary distributions.The paper compares this sample-quality strategy with other reported approaches in Appendix B.9.Source: §3.6 · p. 5

4.1. Implementation Notes of Dirichlet Diffusion Score Model

• Although Jacobi diffusion sampling is more expensive than Gaussian-stationary SDE sampling, precomputing diffused endpoint samples and score gradients permits efficient training with little additional overhead.Only samples starting from 0 and 1 are needed for categorical data.Source: §4.1 · p. 5

4.2. Application to Binarized MNIST

• DDSM obtained competitive performance on the binarized MNIST benchmark.The benchmark results are reported in Table 1.Source: §4.2 · p. 6, Table 1 · p. 6

4.3. Sudoku Generation as a Constraint Satisfying Generation Test

• DDSM with time-dilation achieved up to 100% correctly filled generated Sudoku, compared with 0.31% for the training-data heuristic, and outperformed tested diffusion baselines.Evaluation gave credit only for completely correct puzzles; more sampling steps and time-dilation improved generation accuracy.Source: §4.3 · p. 6, Table 2 · p. 6

4.4. Solving Sudoku via Conditional Generation

• With 128x time-dilation, one sample solved 99.4% of easy and 42.4% of hard Sudoku puzzles; multiple samples solved all tested puzzles.The model required no additional training for solving, and neither Bit Diffusion nor D3PM-uniform solved hard puzzles.Source: §4.4 · pp. 6–7, Table 2 · p. 6

4.5. Generation of Promoter DNA Sequences

• Promoter Designer conditions generation on transcription initiation signal profiles using a dataset of 100,000 promoters, each 1,024 basepairs long.The profiles correspond to promoters spanning a range of activity levels and are obtained from CAGE experiments.Source: §4.5 · p. 7

4.6. Evaluation of Designed Promoter DNA Sequences

• ≤1.32 bits per basepair was the conditional NLL estimate for top-10,000 test-chromosome promoters, versus 1.92 bits for position-specific base-composition baselines.Higher-activity promoters had better NLL estimates; generated sequences were diverse and returned no BLAST hits against the human genome.Source: §4.6 · p. 7
• Designed promoters had genome-like nucleotide composition and motif-location distributions, and comparable predicted activity for high- and low-activity targets.Time-dilation further increased predicted promoter activity.Source: §4.6 · p. 7, Figure 4 · p. 8
• DDSM outperformed baseline discrete-diffusion models on SP-MSE promoter evaluation under a shared Promoter Designer architecture and early-stopping criterion.SP-MSE measures the difference between predicted promoter activity for generated and human-genome sequences; lower is better.Source: §4.6 · p. 7, Table 3 · p. 7

Discrete-data diffusion

• Discrete-data diffusion work includes discrete-time methods [Hoogeboom et al., 2021b; Austin et al., 2021], continuous-time Markov-chain approaches [Sun et al., 2022], and continuous-space encodings with quantization [Chen et al., 2022b; Li et al., 2022].The passage describes quantization-based methods as exceptions to the absence of continuous-time SDE formulations for discrete data.Source: §5 · p. 8

Continuous-space diffusion alternatives

• Concurrent work used Cox-Ingersoll-Ross processes with Gamma stationarity, while reflected SDEs offer another way to restrict score-based diffusion to the probability simplex.The Cox-Ingersoll-Ross approach normalizes a vector of independent processes to unit sum.Source: §5 · p. 9

Latent diffusion

• Latent diffusion can model discrete data through continuous latent variables, whereas this paper describes its latent-to-discrete relationship as fixed rather than learned (Vahdat et al., 2021; Lovelace et al., 2022).Source: §5 · p. 9

Biological sequence design

• Earlier biological sequence design used deep generative models for DNA and structure-conditioned protein sequences, but the authors report no prior human promoter design method or diffusion model for DNA sequence design to their knowledge.Prior work had applied diffusion to protein structure and jointly generated protein structure and sequence (Trippe et al., 2022; Anand & Achim, 2022).Source: §5 · p. 9

Contribution in context

• The paper adds continuous-time SDE diffusion with Dirichlet stationarity to the score-based toolkit and applies it to human promoter design and Sudoku generation.It also proposes time-dilation to improve sample quality.Source: §5 · p. 9

6. Discussion

• DDSM provides a continuous-time diffusion framework for discrete data and a plug-in alternative to Gaussian-stationary SDEs, with biological sequences among its intended applications.The discussion also identifies Sudoku constraint satisfaction and human promoter design as promising results motivating further development.Source: §6 · p. 9
• The continuous diffusion state grows linearly with category count, making the approach best suited to moderate-category data such as DNA, RNA, and protein sequences.The discussion notes memory and computational constraints at very high category counts and suggests bit or hierarchical encoding as potential remedies.Source: §6 · p. 9

A.1. Transition Density Function of Jacobi Diffusion Process

• The Jacobi transition density is expressed through a spectral expansion involving a Beta density and modified Jacobi polynomials.Source: §A.1 · p. 12

A.2. Choosing Speed Factor s in Jacobi Diffusion Processes for Dirichlet Diffusion Score Model

• The speed factor s changes convergence speed without changing the stationary distribution; the paper uses s = 1 for Sudoku and s = 2/(a+b) for promoter design.The two choices target uniform convergence in x space or v space, respectively, and the authors recommend comparing them for specific applications.Source: §A.2 · pp. 12–13

A.3. Weighting Function for Score Matching Loss Invariant to Change-of-Variable

• The proposed weighted score-matching loss is invariant under differentiable change of variables, unlike the unweighted loss discussed in the derivation.Source: §A.3 · p. 13

A.4. Improving Sampling Efficiency for Data with High Number of Categories

• The fast forward sampler gives an exact k−1-fold speedup and lower precomputation memory, with 2.08 vs 2.87 ms for fast vs regular sampling at 10,000 dimensions.Source: §A.4 · p. 13
• Fast sampling produces nearly identical samples to the original multivariate Jacobi process, with no noticeable quality or likelihood decrease reported and perfect Sudoku-generation accuracy.On synthetic data, the fast-sampling model obtained 2.08 bits/dim against a 2 bits/dim ground-truth optimum.Source: §A.4 · p. 13

A.5. Improving Sample Quality by Biasing Reverse Diffusion Toward High-Density Areas

• Time dilation modifies reverse diffusion without changing the score model, biasing samples toward high-density regions that often correspond to higher-quality samples.The authors report that it need not increase time complexity unless the number of sampling steps is increased, and recommend scaling steps with the dilation factor.Source: §A.5 · p. 14

B.1. Implementation of the Jacobi Diffusion Transition Density Function

• For Jacobi density scores, 1000 terms suffice at t = 0.001, 100 at t = 0.01, and 20 at t = 0.1; the experiments conservatively used 1000 terms.The implementation computes Jacobi polynomials by dynamic programming, with time complexity linear in the number of expansion terms.Source: §B.1 · p. 14, Table 5 · p. 15

B.2. Efficient Sampling and Score Computation from Jacobi Diffusion Processes

• Jacobi sampling and score computations can be precomputed for training; generating a sample and its score took 2.87 ms combined for 10,000 dimensions.Presampled dictionaries can be reused for applications with the same forward diffusion processes.Source: §B.2 · p. 15

B.3. Importance Sampling of Time During Denoising Score-Matching Training

• Importance sampling selects high-score time points more often and down-weights their losses to reduce gradient variance during likelihood-weighted training.Source: §B.3 · p. 15

B.4–B.7. Score Model Design, Sampling, Minimum Time, and Stick-Breaking Order

• Implementation choices include Gaussian Fourier time embeddings, Euler–Maruyama sampling, a positive minimum diffusion time, and optional randomization of stick-breaking order.The minimum time avoids the undefined score at t = 0 and numerical issues at very small t; 0.01 or 0.001 is sufficient with typical parameters.Source: §B.4 · p. 16, §B.5 · p. 16, §B.6 · p. 16

B.8. The empirical assessment on the tightness of ELBO

• The ELBO gap is judged tight enough for most applications and can be reduced by lowering t̃0, a time close to zero.Source: Table 6 · p. 16, §B.8 · p. 16

B.9. Comparison of Time Dilation Approach with Other Improving Sample Techniques

• Time dilation consistently outperformed predictor-corrector sampling on Sudoku at the same number of evaluations, although the methods target biased and unbiased sampling, respectively.Source: §B.9 · p. 16, Table 7 · p. 17

C. Application Details; C.1. Binarized MNIST

• The binarized-MNIST application uses a Ho et al. architecture with continuous-time Gaussian Fourier embeddings and univariate Jacobi diffusion.Source: C. Application Details · p. 17, §C.1 · p. 17

C.2. Sudoku Generation and Solving

• Sudoku training uses fully filled grids from a generation heuristic with 0.31% accuracy; the model uses a 20-block transformer and Euler–Maruyama sampling with or without time dilation.The reported Sudoku setup also uses 3200 steps for single-sample accuracy estimates and 200 steps with 8× time dilation for repeated-sample solving.Source: §C.2 · p. 17

C.3. Sudoku performance comparison with baseline methods

• Table 8 covers single-sample Sudoku generation and solving accuracies compared across baseline diffusion methods using the Sudoku transformer architecture.Source: Table 8 · p. 17

C.4. Human Promoter Sequence Design

• The promoter dataset contains 100,000 highly expressed 1024-bp human sequences centered on annotated transcription start sites, split by chromosome for training, validation, and test.Training sequences and transcription profiles receive synchronized random shifts of up to ±100 bp.Source: §C.4 · p. 18
• Promoter Designer conditions sequence generation on CAGE transcription-initiation profiles, and generated sequences are evaluated against test-set promoters and a Sei-model activity score.The baseline diffusion models use the same architecture and validation-set early-stopping criterion.Source: §C.4 · p. 18

• Stronger time-dilation produced more stereotypical digits and overrepresented the common digits “1” and “7”.The authors suggest this exaggerates the dataset’s non-uniform digit distribution.Source: D. Binarized MNIST Generation Examples… · p. 19
• Starting 8x time-dilation at 25% of reverse-diffusion time preserved increased digit quality without biasing toward overrepresented digits.The illustrated sampler used 275 steps.Source: D. Binarized MNIST Generation Examples… · pp. 19, 21
• The sample examples compare Euler–Maruyama runs at 1x, 2x, 4x, and 8x time-dilation, using 100, 200, 400, and 800 steps, respectively.Source: Figure D.5 · p. 19, Figure D.6 · p. 20, Figure D.7 · p. 20

E. Sudoku Transformer Model Performance on the Hard Sudoku Dataset

• Multiple samples solved 100% of easy and hard puzzles (means 1.19 and 753 samples); one 128x-dilated sample solved 99.4% easy and 42.4% hard.Source: E. Sudoku Transformer Model… · p. 22
• The reported results are not directly comparable with supervised baselines because the model was trained on generator-produced full grids, not Sudoku puzzles.Source: E. Sudoku Transformer Model… · p. 22
• The current setup uses 8x time-dilation and 200 steps with a simple Euler–Maruyama sampler, and the authors say it is not optimized.Source: E. Sudoku Transformer Model… · p. 22
• The hard-set experiment tracks required samples by clue count, including 17-clue puzzles, but optimizes total solving time rather than sample count.Source: Figure E.1 · p. 22, E. Sudoku Transformer Model… · p. 22

• The task is to generate promoter sequences matching a target transcription-initiation profile, for which the authors report no prior method.Source: §F.1 · p. 23
• Designed sequences resembled natural promoters in position-specific nucleotide composition and known motif locations and frequencies; predicted activity also tracked expression groups.Source: §F.1 · p. 23
• Known motifs are not the only important sequences, and the rules governing transcription remain incompletely understood.Source: §F.1 · p. 23
• Time-dilation slightly increased Sei-predicted promoter activity and known motif frequencies, while making nucleotide composition slightly more biased.Source: Figure F.1 · p. 24, Figure F.2 · p. 25, Figure F.3 · p. 26