PaperlayerView exact revision on arXiv

Source-linked AI summary

Dirichlet Flow Matching with Applications to DNA Sequence Design

Hannes Stark, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, Tommi Jaakkola

arXiv:2402.05841v2q-bio.BMcs.LG

TL;DR

Flow matching had not yet addressed discrete categorical data, despite its relevance to sequence design. The paper develops Dirichlet flow matching on the simplex, with guidance and distillation, and reports strong DNA-design results, including one-step generation with multiple orders of magnitude speedups over autoregressive models.

  • Problem

    Existing flow-matching formulations focused on continuous spaces and had not treated discrete categorical data, despite applications such as text generation and biological sequence design.

    Source: §1 · p. 1

  • Method

    Dirichlet flow matching transports between uniform simplex noise and vertex-supported data using conditional Dirichlet paths, with score–flow connections for guidance and a sampling process that can be distilled.

    Source: §1 · pp. 1–2, §5 · p. 9

  • Results

    Dirichlet flow matching outperforms linear flow matching and multiple discrete diffusion approaches on regulatory DNA tasks; distillation enables one-step generation with multiple orders of magnitude speedups and marginal performance impact.

    Source: §5 · p. 9

  • Takeaways & Limitations

    The results support Dirichlet flow matching as a flexible framework for biological sequence generation, with guidance and distillation and strong performance on the evaluated tasks.

    Source: §5 · p. 9

Abstract

from arXiv · show

Discrete diffusion or flow models could enable faster and more controllable sequence generation than autoregressive models. We show that na\"ive linear flow matching on the simplex is insufficient toward this goal since it suffers from discontinuities in the training target and further pathologies. To overcome this, we develop Dirichlet flow matching on the simplex based on mixtures of Dirichlet distributions as probability paths. In this framework, we derive a connection between the mixtures' scores and the flow's vector field that allows for classifier and classifier-free guidance. Further, we provide distilled Dirichlet flow matching, which enables one-step sequence generation with minimal performance hits, resulting in $O(L)$ speedups compared to autoregressive models. On complex DNA sequence generation tasks, we demonstrate superior performance compared to all baselines in distributional metrics and in achieving desired design targets for generated sequences. Finally, we show that our classifier-free guidance approach improves unconditional generation and is effective for generating DNA that satisfies design targets. Code is available at https://github.com/HannesStark/dirichlet-flow-matching.

1. Introduction

Discrete categorical generation remains an important application gap for flow matching, including biological sequence design. The paper introduces Dirichlet flow matching on the simplex, develops guidance and distillation, and reports strong results on DNA design tasks.

1. Introduction

  • Flow matching offers simulation-free training of continuous normalizing flows between noise and data distributions.Its flexible iterative noising processes have been applied to non-Euclidean spaces, but existing formulations had not treated discrete categorical data.Source: §1 · p. 1
  • Dirichlet flow matching models discrete categorical data as transport on the simplex, using conditional noising that produces time-evolving Dirichlet distributions.The transport runs from a uniform density over the simplex to a finitely supported distribution over its vertices.Source: §1 · p. 1
  • Linear flow matching on the simplex suffers from pathological behavior caused by contracting support, whereas Dirichlet flow matching is designed to avoid these shortcomings.The paper reports that Dirichlet flow matching is simpler computationally and conceptually than DDSM.Source: §1 · p. 1
  • A connection between mixture-of-Dirichlet scores and the resulting flow enables classifier guidance and classifier-free guidance.The sampling process defines a deterministic map between noise and data, which the authors use to distill generation to a single forward pass.Source: §1 · pp. 1–2

5. Conclusion

  • Dirichlet flow matching outperforms linear flow matching and multiple discrete diffusion approaches on regulatory DNA design tasks, while distillation enables one-step generation with multiple orders of magnitude speedups.Distillation only marginally impacts performance; guided generation also designs cell-type-specific enhancers.Source: §1 · p. 2

1. Introduction

  • The evaluation reports improved promoter design, enhancer distributional similarity, guided DNA design, and one-step generation with minimal performance degradation.On enhancer datasets, FBD is 1.9 vs. 36.0 in Melanoma DNA and 1.0 vs. 25.2 in Fly Brain DNA, compared with autoregressive language modeling.Source: §1 · p. 2

2. Background

The background introduces flow matching as transport from noise to data, reviews major discrete-diffusion approaches, and explains the promoter and enhancer design setting.

2.1. Flow Matching

  • Flow matching regresses a neural network against a vector field transporting noise samples to data, with a conditional probability path from noise toward the data distribution.The path and associated vector field can be defined through a conditional flow map that transports a noise sample to an intermediate distribution.Source: §2.1 · p. 2

2.2. Related Discrete Diffusion Model Works

  • Existing discrete diffusion methods include simplex-based, unconstrained continuous-space, discrete-noise, and latent-discrete approaches.DDSM is a related simplex-based method using Jacobi diffusion and stick-breaking; the reviewed simplex-based approaches do not use Dirichlet distributions as intermediate noising distributions.Source: §2.2 · p. 3

2.3. Promoter and Enhancer DNA

  • Promoters and enhancers are noncoding DNA regions that regulate gene expression; promoters lie next to genes, while enhancers can be distant in sequence and often act in specific cell types.Cell-type specificity enables targeting, such as targeting cancer cells.Source: §2.3 · p. 3
  • Prior enhancer designs used optimization guided by a cell-type activity classifier, while this work aims to learn and sample conditional data distributions that better capture empirical enhancer patterns.The authors frame this as a foundation for more principled conditional sequence design.Source: §2.3 · p. 3

3. Method

The method transports data on the simplex using Dirichlet probability paths, avoiding the shrinking support and discontinuous fields of linear flow matching. It also derives score-to-flow relationships for guidance and distills deterministic sampling into a single step.

3.1. Flow Matching on the Simplex

  • The simplex flow transports a uniform prior to categorical data represented as a distribution over simplex vertices, using a denoising classifier to parameterize the vector field.The cross-entropy and flow-matching losses have the same minimizer; this objective keeps the learned field tangent to the simplex and avoids evaluating the conditional field during training.Source: §3.1 · p. 3

3.2. Designing Simplex Flow Matching

  • Linear flow matching’s posterior support contracts over time, forcing possible destination vertices to be eliminated and making the field increasingly discontinuous.For t > 1/k, Proposition 1 bounds posterior support to at most k − 1 vertices; the authors identify these pathologies as especially harmful at higher dimensions.Source: §3.2 · p. 4

3.3. Dirichlet Flow Matching

  • Dirichlet FM uses a conditional path that retains support over the entire simplex, with a vector field designed to generate that path.The field points toward the target vertex and is rescaled by a factor that depends on the target coordinate and time.Source: §3.3 · p. 5

3.4. Guidance

  • A derived linear relationship between Dirichlet-mixture scores and marginal flows supports classifier and classifier-free guidance.Classifier-free guidance combines conditional and unconditional scores; guidance strengths γ > 1 can enhance efficacy.Source: §3.4 · pp. 5–6
  • Dirichlet FM scales to larger categorical dimensions much better than linear FM in toy experiments measured by KL divergence to the target distribution.The comparison uses 512k generated samples for increasing category counts K.Source: Figure 4 · p. 6

3.5. Distillation

  • Because Dirichlet FM sampling is a deterministic ODE map, a student can be trained to reproduce the teacher distribution in one step.The experiments use a 100-step teacher and train the student on teacher-generated noise–target pairs.Source: §3.5 · p. 6

4. Experiments

Experiments test Dirichlet flow matching on simplex distributions, promoter design, and enhancer generation. Across these tasks, the method performs well against baselines, while guidance improves conditional and unconditional enhancer generation.

4.1. Simplex dimension toy experiment

  • Dirichlet FM overfits the toy categorical distribution regardless of dimension, whereas linear FM cannot overfit as the number of categories increases.The comparison evaluates KL divergence between the empirical distribution of 512,000 model samples and the target distribution.Source: §4.1 · pp. 6–7

4.2. Promoter DNA sequence design

  • A single-step distilled Dirichlet FM retains almost the same promoter-design performance and outperforms the other methods, with 100× fewer function evaluations than diffusion models and 1,024× fewer than the language model.Dirichlet FM improves upon linear FM; only those two methods outperform the language-model baseline.Source: §4.2 · p. 7, Table 1 · p. 7

4.3. Enhancer DNA design

  • Dirichlet FM outperforms the language model on both enhancer datasets and linear FM on human melanoma; distillation minimally affects FBD while requiring one step.The evaluation compares generated distributions with the unconditional data distribution using FBD.Source: Table 2 · p. 7, §4.3 · p. 8
  • For conditional enhancer design, increasing γ can raise target-class probability close to 100%; FBD improves to a minimum, then worsens as diversity decreases.For PNG, guided sequences have better FBD than the available functional sequences of Taskiran et al. (2023); classifier-guidance improvements are significant but smaller.Source: §4.3 · p. 8, Figure 5 · p. 8

Q3: Can guidance improve unconditional generation?

  • Sampling a class by empirical frequency and guiding toward it significantly improves unconditional enhancer sample quality over unguided Dirichlet FM and baselines.FBD first improves with increased guidance, then worsens as increased guidance deteriorates sample diversity.Source: Q3: Can guidance improve unconditional… · pp. 8–9, Figure 6 · p. 8

5. Conclusion

Dirichlet flow matching addresses linear flow matching’s simplex pathologies and supports guidance and distillation. Across three DNA datasets, it performs strongly on sequence design tasks, while distillation enables one-step generation with marginal performance impact.

5. Conclusion

  • Dirichlet flow matching models discrete data through a generative process on the simplex, avoiding the pathological behavior identified in linear flow matching.Its marginal flow connects to mixture-of-Dirichlet scores, supporting guidance; its deterministic sampling process also permits distillation.Source: §5 · p. 9
  • Across three datasets, Dirichlet flow matching outperforms linear flow matching and multiple discrete diffusion approaches on regulatory DNA design tasks.Distillation enables one-step generation with only a marginal performance impact and multiple-orders-of-magnitude speedups over autoregressive models for long sequences.Source: §5 · p. 9
  • Guided Dirichlet flow matching generates cell-type-specific enhancers, and the framework is presented as a promising direction for further discrete-data applications.The conclusion describes its performance for biological sequence generation as strong.Source: §5 · p. 9

Impact Statement

The paper applies an iterative generative process for discrete data to controllable regulatory DNA design and describes potential societal impacts.

Impact Statement

  • The paper presents an iterative generative modeling process for discrete data and applies it to controllable sequence generation.Its application is regulatory DNA design.Source: Impact Statement · p. 9
  • The stated potential use cases include discovering biological regulatory mechanisms, drug development, gene therapy, and other therapies.The authors characterize current societal impacts of this application as overwhelmingly positive.Source: Impact Statement · p. 9
  • The impact statement says that none of the framework’s many possible societal impacts need specific highlighting.This statement accompanies the application to regulatory DNA design.Source: Impact Statement · p. 9

A. Method Details

The appendix characterizes a support pathology of linear flow matching and derives a conditional Dirichlet vector field by matching probability flux to the changing mass of a simplex region.

Proof of Proposition 1

  • For t>1/k, the converged linear-flow posterior has support over at most k−1 vertices for x sampled from p_t.This is the stated result for k=2,...,K.Source: Proof of Proposition 1 · p. 13

A.1. Dirichlet Conditional Vector Field

  • The conditional field points toward target vertex e_i, with magnitude depending only on x_i and t; its scale is set by matching boundary flux to the changing mass of a simplex region.The construction equates the probability flux across the hyperplane with the rate of change of the probability in the region it bounds.Source: §A.1 · pp. 13–14
  • The mass of S1 is calculated using an incomplete beta function, and its first-parameter derivative gives the rate of change of that mass.The regularized incomplete beta function is identified as the Beta-distribution CDF; the derivative is not available in closed form.Source: §A.1 · p. 14
  • The derivation checks the transport equation by comparing the divergence of the probability flux with the time derivative of the density.The appendix states that the first and second terms cancel in this check.Source: §A.1 · p. 15

B.1. Training and Inference

Training uses one-hot sequences and time-conditioned per-position probability predictions. Inference starts from simplex noise and integrates the learned field over a specified number of steps.

B.1. Training and Inference

  • Training supplies one-hot sequences to a neural network that predicts probabilities at each sequence position, samples exponential times, and optimizes a prediction loss.Source: §B.1 · p. 16
  • Inference initializes each sequence position from a uniform Dirichlet distribution and advances the sequence through N integration steps up to t_max.The network predicts probabilities at each step and position.Source: §B.1 · p. 16
  • At each position, inference computes conditional vector fields across simplex dimensions and then calculates the marginal vector field.Source: §B.1 · p. 16

N · ut(xt,l) end for end for return xt

This appendix describes the training and inference setups, model architectures, and guidance procedures used in the experiments. It also reports implementation and computational details.

Promoter Design

Enhancer Design

Classifier-free Guidance

  • Classifier-free guidance uses class labels on 0.7 of training examples, combines conditioned and unconditioned probability predictions, and projects the result onto the simplex.Projection is needed because guidance factors γ > 1 can produce negative values.Source: N · ut(xt,l) end for end for return xt · p. 17

Classifier Guidance

  • Classifier guidance converts model probabilities to scores, adds unconditional and conditional scores, converts back to probabilities, and projects them onto the simplex before integration.The noisy classifier is trained with early stopping on accuracy for approximately 800 epochs; gradients are obtained by automatic differentiation.Source: N · ut(xt,l) end for end for return xt · p. 17

Implementation

  • Distillation trains the student on teacher-generated noise-target pairs across approximately 6 billion sequences; enhancer training takes seven hours for 200 epochs on RTX A600 GPUs.The largest model uses approximately 8GB of RAM during training.Source: N · ut(xt,l) end for end for return xt · p. 17

Architecture

  • Models use 20 layers of 1D convolutions with time embeddings and normalization; promoter experiments use the DDSM architecture, while other experiments replace group norm with layer norm.Transformer architectures performed worse; the Fly Brain classifier has five layers.Source: N · ut(xt,l) end for end for return xt · p. 17

B.2. Data

The enhancer datasets are ATAC-seq accessibility data, with stated train/validation/test splits for fly-brain and human melanoma sequences. Accessibility is not sufficient to establish enhancer activity.

Data

  • ATAC-seq measures chromatin accessibility, a necessary but insufficient condition for a sequence to act as an enhancer.Thus, the datasets do not guarantee that all sequences enhance transcription in a cell type.Source: §B.2 · p. 18
  • The fly-brain dataset contains 104,665 sequences split 83,726/10,505/10,434 for train/validation/test, following Taskiran et al. (2023).Source: §B.2 · p. 18
  • The human melanoma dataset contains 88,870 sequences split 70,892/8,966/9,012 for train/validation/test, and has little class imbalance.Source: §B.2 · p. 18

C. Additional Results

An analytically tractable toy experiment tests classifier guidance by comparing generated classes with the known conditional distribution. Its generated distribution closely matches the ground truth, and classifier performance plots closely match prior reports.

Analytical toy experiment for classifier guidance

  • The toy setup uses a binary conditioning variable and a 20-class categorical distribution, with conditional probabilities, scores, and the Dirichlet FM vector field computed analytically.Source: C. Additional Results · p. 18
  • The generated class distribution closely matches the ground-truth probabilities, empirically confirming the classifier-guidance procedure’s probability-to-score conversion and linear-system inversion.Source: C. Additional Results · p. 18, Figure 10 · p. 19

Additional classifier results

  • Classifier performance scatter plots for fly-brain and human-melanoma classes closely match the results of Atak et al. (2021) and Janssens et al. (2022).The plots compare area under the ROC curve with area under the precision-recall curve.Source: Figure 11 · p. 19
Loading 2402.05841v2…