[Paper Review] Dirichlet Diffusion Score Model for Biological Sequence Generation

Pavel Avdeyev, Chenlai Shi+3 more|PubMed|May 18, 2023

Genetic and phenotypic traits in livestock11 references8 citations

TL;DR

A Dirichlet diffusion score model (DDSM) extends score-based diffusion to discrete data by diffusion on the probability simplex with Dirichlet stationary distribution, enabling constrained sequence design such as Sudoku puzzles and human promoter DNA sequences.

ABSTRACT

Designing biological sequences is an important challenge that requires satisfying complex constraints and thus is a natural problem to address with deep generative modeling. Diffusion generative models have achieved considerable success in many applications. Score-based generative stochastic differential equations (SDE) model is a continuous-time diffusion model framework that enjoys many benefits, but the originally proposed SDEs are not naturally designed for modeling discrete data. To develop generative SDE models for discrete data such as biological sequences, here we introduce a diffusion process defined in the probability simplex space with stationary distribution being the Dirichlet distribution. This makes diffusion in continuous space natural for modeling discrete data. We refer to this approach as Dirchlet diffusion score model. We demonstrate that this technique can generate samples that satisfy hard constraints using a Sudoku generation task. This generative model can also solve Sudoku, including hard puzzles, without additional training. Finally, we applied this approach to develop the first human promoter DNA sequence design model and showed that designed sequences share similar properties with natural promoter sequences.

Motivation & Objective

• Motivate the need for generative models that satisfy hard constraints in discrete biological sequence design.
• Introduce a continuous-time diffusion process on the probability simplex with Dirichlet stationary distribution.
• Develop a score-based diffusion framework for k-category discrete data via stick-breaking and Jacobi diffusions.
• Demonstrate the model on constrained data tasks (Sudoku) and a real-world biological application (human promoter design).
• Provide training, sampling, and likelihood estimation tools for Dirichlet diffusion in discrete spaces.

Proposed method

• Construct forward diffusion in probability simplex using multivariate Jacobi diffusions that converge to Dirichlet(…) stationary distributions via stick-breaking.
• Define reverse-time SDE and probability flow ODE using learned score function to sample discrete data.
• Train with a generalized score-matching loss tailored for Jacobi diffusion, including a likelihood-inspired weighting and change-of-variable considerations.
• Enable efficient sampling by leveraging two-point initial conditions (0/1) and precomputed densities, with optional time-dilation to improve sample quality.
• Provide a likelihood framework via probability flow ODE and an ELBO-based bound for discrete data likelihood.
• Apply a conditional generation setup for promoter design by conditioning on transcription initiation signal profiles.

Experimental results

Research questions

• RQ1 Can continuous-time diffusion on the probability simplex, with Dirichlet stationarity, effectively model discrete data such as biological sequences?
• RQ2 How can score-based diffusion be adapted for multi-category discrete data via stick-breaking and Jacobi diffusion processes?
• RQ3 To what extent can DDSM generate highly constrained discrete data (e.g., Sudoku puzzles) and design real-world biological sequences (promoter DNA) while preserving target properties?
• RQ4 What are practical training, sampling, and likelihood estimation strategies for Dirichlet diffusion in discrete spaces?

Key findings

• DDSM achieves competitive negative log-likelihood on binarized MNIST compared to several baselines.
• The model can generate and solve Sudoku puzzles, achieving up to 100% generation accuracy with time-dilation and conditioning techniques.
• A first method for designing human promoter DNA sequences is demonstrated, with generated promoters sharing properties with natural promoters and achieving competitive promoter activity predictions.
• Time-dilation and sampling strategies improve sample quality and puzzle-solving success, outperforming several diffusion baselines in constrained generation tasks.
• The conditional promoter design setup enables controlled generation conditioned on transcription initiation profiles, yielding sequences with comparable activity to ground-truth promoters.