Dirichlet Flow Matching with Applications to DNA Sequence Design
This work addresses the need for faster and more controllable sequence generation in bioinformatics, particularly for DNA design, though it is incremental as it builds on existing flow matching methods.
The paper tackled the problem of generating DNA sequences more efficiently and controllably than autoregressive models by developing Dirichlet flow matching on the simplex, which avoids discontinuities and enables one-step generation with O(L) speedups, achieving superior performance on distributional metrics and design targets.
Discrete diffusion or flow models could enable faster and more controllable sequence generation than autoregressive models. We show that naïve 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.