Dirichlet Flow Matching with Applications to DNA Sequence Design
Hannes Stark Bowen Jing Chenyu Wang Gabriele Corso Bonnie Berger Regina Barzilay Tommi Jaakkola
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.