Dirichlet Flow Matching with Applications to DNA Sequence Design

Introduction

Flow matching (FM) is a generative modeling framework that provides a simulation-free means of training continuous normalizing flows (CNFs) between noise and data distributions (Lipman et al., 2022; Liu et al., 2022; Albergo & Vanden-Eijnden, 2022). Such flow models can be viewed as generalizing diffusion models (Song et al., 2021) to permit more flexible design of iterative noising processes, which have proven useful in generative modeling on non-euclidean spaces such as compact Riemannian manifolds (Chen & Lipman, 2023; Bose et al., 2023). However, existing formulations of flow matching have focused on continuous spaces and have yet to treat discrete categorical data—a notable shortcoming considering the many important applications of discrete generative modeling such as text generation and biological sequence design (Avdeyev et al., 2023).

We evaluate Dirichlet FM on three DNA sequence datasets with several complex tasks that pose diverse challenges. First, we demonstrate that Dirichlet FM better generates promoter DNA sequences with desired regulatory activity compared to baselines such as an autoregressive language model or discrete diffusion models, including DDSM (Avdeyev et al., 2023). Second, we show on two enhancer DNA sequence datasets that Dirichlet FM improves upon autoregressive language modeling in capturing the data distribution with an FBD (distributional similarity) of 1.9 vs. 36.0 in Melanoma DNA and 1.0 vs. 25.2 in Fly Brain DNA. Third, we demonstrate that Dirichlet FM guidance can improve unconditional sequence generation and generate cell-type specific enhancer DNA sequences that improve upon the (experimentally validated) sequences of Taskiran et al. (2023). Lastly, distilled Dirichlet FM generates sequences in a single step, resulting in orders-of-magnitude speedups with minimal performance degradation.

Background

In flow matching (Lipman et al., 2022; Liu et al., 2022; Albergo & Vanden-Eijnden, 2022), we consider a noisy distribution $\mathbf{x}_{0}\sim q_{0}$ and data distribution $\mathbf{x}_{1}\sim p_{\text{data}}$ and regress a neural network against a vector field that transports $q(\mathbf{x}_{0})$ to $p_{\text{data}}(\mathbf{x}_{1})$. To do so, we define a conditional probability path—a time-evolving distribution $p_{t}(\mathbf{x}\mid\mathbf{x}_{1}),t\in$ conditioned on $\mathbf{x}_{1}$ with boundary conditions $p_{0}(\mathbf{x}\mid\mathbf{x}_{1})=q(\mathbf{x})$ and $p_{1}(\mathbf{x}\mid\mathbf{x}_{1})\approx\delta(\mathbf{x}-\mathbf{x}_{1})$. We additionally assume knowledge of a conditional vector field $u_{t}(\mathbf{x}~{}\mid~{}\mathbf{x}_{1})$ that generates $p_{t}(\mathbf{x}\mid\mathbf{x}_{1})$, i.e. satisfying the transport equation

interpolates between noise $p_{0}=q_{0}$ and data $p_{1}\approx p_{\text{data}}$ and is generated by the marginal vector field

Thus, by learning and integrating a neural network $\hat{v}(\mathbf{x},t;\theta)\approx v_{t}(\mathbf{x})$, we can generate data from noisy samples $\mathbf{x}_{0}\sim q$. The core design decision is the choice of appropriate conditional probability path $p_{t}(\mathbf{x}\mid\mathbf{x}_{1})$ and associated vector field $u_{t}(\mathbf{x}\mid\mathbf{x}_{1})$. Although it is possible to define these directly, it is often simpler to instead define a conditional flow map $\psi_{t}(\mathbf{x}_{0}\mid\mathbf{x}_{1})$ that directly transports $\mathbf{x}_{0}\sim q$ to the intermediate distribution $p_{t}(\mathbf{x}\mid\mathbf{x}_{1})$. The flow map immediately provides the corresponding vector field:

With this formulation, the required boundary conditions simplify to $\psi_{0}(\mathbf{x}_{0}\mid\mathbf{x}_{1})=\mathbf{x}$ and $\psi_{1}(\mathbf{x}_{0}\mid\mathbf{x}_{1})=\mathbf{x}_{1}$. As advocated by several works (Liu et al., 2022; Lipman et al., 2022; Pooladian et al., 2023; Tong et al., 2023), the flow map (also called interpolant) is often chosen to follow the simplest possible path between the two endpoints—e.g., linear in Euclidean spaces and geodesic on Riemannian manifolds (Chen & Lipman, 2023).

2 Related Discrete Diffusion Model Works

Existing discrete diffusion frameworks can be split into 4 categories. Firstly, simplex-based approaches frame discrete data as vertices of a simplex and generate it starting from a Dirichlet prior over the whole simplex (Richemond et al., 2022; Floto et al., 2023). Among those, DDSM (Avdeyev et al., 2023) is most related to our work and converges to a Dirichlet distribution via Jacobi diffusion processes and the stick-breaking transform. We note that none of these simplex-based approaches feature Dirichlet distributions as intermediate distributions of the noising process—a key aspect of our approach.

The second class of discrete diffusion models fully relaxes discrete data into continuous space without any constraints and uses, e.g., a standard Gaussian as prior (Han et al., 2022; Chen et al., 2023; Frey et al., 2024). The third paradigm, established by D3PM (Austin et al., 2021), operates on discrete samples of noise distributions constructed by injecting discrete noise into data (Campbell et al., 2022; Igashov et al., 2024; Vignac et al., 2023; Penzar et al., 2023). Lastly, latent discrete diffusion models train an additional network to obtain continuous latents for which they train a conventional diffusion model (Dieleman et al., 2022; Li et al., 2023).

3 Promoter and Enhancer DNA

DNA is a sequence with base pairs as tokens (3 billion for humans) and a vocabulary of 4 nucleotides (A, T, C, G). Parts of DNA encode genes that are transcribed into mRNA and then translated to functional proteins. Promoters and enhancers refer to noncoding portions of DNA that regulate the expression level of these genes and play important roles in eukaryotic organisms such as humans (Dunham et al., 2012; Luo et al., 2020). More specifically, a promoter for a gene is the DNA sequence next to the gene where the transcriptional machinery binds and starts transcribing DNA to mRNA (Haberle & Stark, 2018). Meanwhile, enhancers are sequences that can be distant in the DNA sequence (millions of base pairs) but are close in 3D space (Panigrahi & O’Malley, 2021) and regulate the recruitment of this transcriptional machinery. Unlike promoters, enhancers often regulate transcription in specific cell types. Hence, while both types of DNA subsequences are important for gene therapy (Whalen, 1994), the cell type specificity of enhancers enables targeting, e.g., only cancer cells.

Designing enhancers. Recently Taskiran et al. (2023) and de Almeida et al. (2023) demonstrated cell type-specific enhancer design via an optimization procedure starting from an initial sequence guided by a cell-type activity classifier. However, such sequence designs may not follow the empirical distribution of enhancers, which would be captured by a generative model. Our work sets the foundation for more principled conditional sequence design by learning and drawing from conditional data distributions.

Method

Let $S_{K}$ be the probability simplex in $K$-dimensional space:

Given a $K$-class categorical distribution with probabilities $p_{i},\sum_{i=1}^{K}p_{i}=1$, we relax this distribution into continuous space by converting it to a mixture of point masses at the vertices of $S_{k}$ (with $\mathbf{e}_{i}$ as the $i$th one-hot vector):

We then define the noisy prior to be the uniform density on the simplex, or a Dirichlet distribution with parameter vector $\boldsymbol{\alpha}$ given by the all ones vector:

Our objective is then to learn a vector field, using some choice of conditional probability path (discussed later), to transport $q_{0}$ to $p_{\text{data}}$. Typically, the neural network directly parameterizes the vector field and is trained via the $L_{2}$-like conditional flow-matching loss

where the expectation is taken over $t\sim\mathcal{U},\mathbf{x}_{1}\sim p_{\text{data}},\mathbf{x}\sim p_{t}(\cdot\mid\mathbf{x}_{1})$. However, we instead train a denoising classifier via a cross-entropy loss

At inference time, we then parameterize the vector field via

It can be shown (Appendix A) that the two losses have the same minimizer, and thus, the cross-entropy is a valid flow-matching objective. The advantages of this approach are twofold: (1) it ensures that the learned vector field is restricted to the tangent plane of the simplex (i.e., the components sum to zero), and (2) the conditional vector field does not need to be evaluated at training time.

For simplicity, our discussion focuses on modeling categorical data on the simplex. However, in practice, we are interested in sequences of variables relaxed onto the multi-simplex $S_{K}^{N}$. At inference time, the simplices depend on each other through a learned denoiser that outputs token-wise logits conditioned on all noisy inputs.

2 Designing Simplex Flow Matching

As outlined in Section 2.1, there are two options to define a conditional probability path $p_{t}(\mathbf{x}\mid\mathbf{x}_{1})$ and corresponding vector field $u_{t}(\mathbf{x}\mid\mathbf{x}_{1})$ to train a flow model with:

Interpolant perspective: Define an interpolant $\psi_{t}(\mathbf{x}_{0}\mid\mathbf{x}_{1})$, which provides the density $p_{t}(\mathbf{x}\mid\mathbf{x}_{1})$ implicitly but allows one to easily sample from it, and obtain the conditional vector field trivially by taking the derivative $u_{t}=\partial_{t}\psi_{t}$.

Probability path perspective: Define $p_{t}(\mathbf{x}\mid\mathbf{x}_{1})$ explicitly and solve for $u_{t}(\mathbf{x}\mid\mathbf{x}_{1})$ that satisfies the transport equation which can be non-trivial.

Following extant works on flow matching, the most natural way to proceed for the simplex would be to follow the interpolant perspective and use the linear flow map employed in Lipman et al. (2022); Pooladian et al. (2023):

Since $S_{K}$ is a Euclidean space, these operations remain well-defined, and the interpolant transports all points on the simplex to $\mathbf{x}_{1}$ at $t=1$ via straight paths. However, a pathological property emerges when such conditional probability paths are marginalized over $p_{\text{data}}(\mathbf{x}_{1})$ over the course of flow matching training:

Suppose that a flow matching model is trained with the linear flow map (Equation 11). Then, for all $k=2,\ldots K$ and $\mathbf{x}\sim p_{t}(\mathbf{x})$, the converged model posterior $p(\mathbf{x}_{1}\mid\mathbf{x})\propto p_{t}(\mathbf{x}\mid\mathbf{x}_{1})p_{\text{data}}(\mathbf{x}_{1})$ has support over at most $k-1$ vertices for times $t>1/k$.

Conceptually, this means that as the model transports samples on the simplex at $t=0$ to the vertices of the simplex at $t=1$, it must eliminate or rule out a possible destination vertex at each of the times $\frac{1}{K},\frac{1}{K-1},\ldots\frac{1}{2}$, if not earlier. As $K$ becomes large, an increasingly large fraction of the model capacity must be allocated to a smaller and smaller fraction of the total time and trajectory length—indeed, for all $K$, the posterior for times $t>1/2$ reduces to the $\arg\max$ operator. Further, since the marginal field is increasingly discontinuous (i.e., rapidly changing directions and settings entries to zero corresponding to eliminated vertices), the model becomes increasingly sensitive to integration step size. We posit—and empirically verify in Section 4—that these factors significantly hurt the performance of linear flow matching, especially for higher dimensionalities $K$.

Upon examination of the marginal probability paths (Figure 2), it becomes clear that this pathological behavior is due to the shrinking support of the conditional paths that arise from linear FM. Unfortunately, by following the interpolant perspective, we are unable to directly control the conditional path. Hence, to obtain a method that does not suffer from linear FM’s pathologies, we proceed with the probability path perspective and directly define the conditional probability path so that it has support on the entire simplex at all times, as described next.

3 Dirichlet Flow Matching

Recall the Dirichlet distribution’s probability density function:

Following the probability path perspective, we first define a conditional probability path with $t\in[0,\infty)$ as:

When $t=0$, this is equal to the uniform noise distribution (Equation 7). As $t\rightarrow\infty$, the $i$th entry of the $\boldsymbol{\alpha}$ parameter vector increases while the others remain constant, concentrating the density towards a point mass on the $i$th vertex, corresponding to the $t=1$ boundary condition in standard flow matching.We continue to call the data sample $\mathbf{x}_{1}$, and in practice, integrate to some large fixed time (typically $t=8$) and take the $\arg\max$ of the final model posterior. Hence, this family of Dirichlet distributions provides a conditional probability path with the required boundary conditions while retaining support over the entire simplex, as desired.

Since we have chosen a conditional probability path directly rather than implicitly via an interpolant, it is more difficult to obtain the corresponding conditional vector field $u_{t}(\mathbf{x}\mid\mathbf{x}_{1})$. Indeed, there is an infinite number of such fields that generate the desired evolution of $p_{t}$. Motivated by the basic form of the linear FM, we generalize it via the following ansatz:

That is, the flow still points directly towards the target vertex $\mathbf{e}_{i}$, but is rescaled by a $x_{i}$-dependent factor. The conditional vector field of linear FM satisfies this form with $C(x_{i},t)=1/(1-t)$ dependent only on $t$; we introduce the additional $x_{i}$-dependence to control the contraction of probability mass towards $\mathbf{e}_{i}$. In Appendix A.1, we derive the $C(x_{i},t)$, which gives rise to Dirichlet probability paths to be:

is a derivative of the regularized incomplete beta function $I_{x}(a,b)$. Figure 3 compares the magnitude of the resulting vector field with that of linear FM as a function of distance from the target vertex. As anticipated, the field vanishes both at the target vertex and on the $(K-2)$-dimensional face directly opposite it instead of monotonically scaling with distance from the target vertex as in linear FM. This second condition means the probability density is never fully drawn away from the face and resolves the pathological behavior of linear FM. The resulting probability paths and vector fields are visualized on the simplex in Figure 1; they are continuous and smooth, unlike in linear FM.

4 Guidance

A key attribute of iterative generative models is the ability to continuously and gradually bias the generative process towards a class label with user-specified strength, a technique known as guidance (Dhariwal & Nichol, 2021; Ho & Salimans, 2022). Initially proposed in the context of diffusion models, where the generative process follows the score $\hat{s}(\mathbf{x},t;\theta)\approx\nabla_{\mathbf{x}}\log p_{t}(\mathbf{x})$ of the noisy data distribution, guidance is implemented by taking a linear combination of the unconditional and conditional score models

with $\gamma>0$ and running the generative process with this adjusted score. In the context of flow matching, Dao et al. (2023); Zheng et al. (2023) derived a relationship between the score and marginal flow for certain Gaussian probability paths and implemented flow guidance by propagating the effects of standard score adjustments to the resulting flow fields. We implement guidance for Dirichlet FM by deriving a similar relationship between the score and flow field, detailed below. Note that when $\gamma=1$, this adjustment precisely mimics the flow that would result from training only on points with class label $y$; however, similar to prior works, we find $\gamma>1$ futher enhances the guidance efficacy.

Relationship between flow and score. For the Dirichlet marginal probability path, the score can be obtained from the model posterior via the denoising score-matching identity (Song & Ermon, 2019):

We can differentiate Equation 14 to obtain a matrix equation

where $\mathbf{D}$ is invertible since it is diagonal with nonnegative entries. Thus, a linear relationship exists between the marginal flow and the score arising from the same model posterior.

Classifier-free guidance. Suppose we have class-conditional and unconditional flow models $\hat{v}(\mathbf{x},t,y;\theta)$ and $\hat{v}(\mathbf{x},t,\varnothing;\theta)$. Since a linear combination of scores results in a linear combination of flows, we similarly implement guidance by integrating

Classifier guidance. In cases where a conditional flow model is unavailable, we use the gradient of a noisy classifier to obtain a conditional score from an unconditional score:

The conditional scores can then be converted into a model posterior $\hat{p}$ and then a marginal flow $\hat{v}$ via Equation 21. However, a direct application is not possible because the classifier gradients do not have the appropriate off-simplex components to ensure a valid model posterior (i.e., $\hat{p}\in S_{K}$) when operated on by $\mathbf{D}^{-1}$. Instead, we modify Equation 21 via

5 Distillation

The aim of distillation (Salimans & Ho, 2022; Song et al., 2023; Yin et al., 2023) is to reduce the inference time of the iterative generative process by reducing the number of steps while retaining sample quality. However, for discrete diffusion models (see Section 2.2) or autoregressive language models, no distillation techniques exist, and it is unclear how to distill generative models based on discrete noise. For Dirichlet FM, inference is a deterministic ODE integration defining a map between the prior and target distribution. Hence, we can distill the teacher model (using 100 steps in our experiments) into a student model representing the map. For this, we sample the teacher to obtain pairs of noise and training targets. We use these to train the student model to reproduce the teacher distribution in a single step. With this, we are able to reduce inference times and provide the first demonstration of distillation for flow matching and for iterative generative models of discrete data.

Experiments

Figure 4 shows these KL-divergences for increasing sizes $K$ of the categorical distribution. Dirichlet FM is able to overfit on the simple distribution regardless of $K$. Meanwhile, Linear FM is unable to overfit on simple categorical distributions as $K$ increases, illustrating the practical implications of the pathological probability paths and discontinuous vector fields as discussed in Section 3.2.

2 Promoter DNA sequence design

We next assess the ability of Dirichlet FM to design DNA promoter sequences conditioned on a desired promoter profile. The experimental setup and evaluation closely follow that of DDSM (Avdeyev et al., 2023).

Task. We train Dirichlet FM conditioned on a profile by providing it as additional input to the vector field. Following Avdeyev et al. (2023), we evaluate generated sequences with the mean squared error (MSE) between their predicted regulatory activity and that of the original sequence corresponding to the input profile. The regulatory activity is determined by the promoter-related predictions of Sei (Chen et al., 2022), a model trained on various regulatory signals.

Baselines. We compare Dirichlet FM with linear FM, discrete diffusion methods, and a language model that autoregressively generates the base pairs. The discrete diffusion baseline most related to our work is the simplex-based DDSM (Avdeyev et al., 2023). Our two other diffusion baselines are Bit Diffusion (Chen et al., 2023) and D3PM (Austin et al., 2021). All methods use the same DNA modeling architecture and training protocol that was designed and tuned by (Avdeyev et al., 2023) for DDSM. See Appendix B.1 for implementation details.

Results. Dirichlet FM is the only method that is able to outperform the language model baseline (Table 1). Linear FM performs the worst, empirically confirming the drawbacks of its discontinuous vector field and pathological probability paths. The second best method in this comparison is the distilled version of Dirichlet FM, which retains almost the same performance. This means that our distilled Dirichlet FM outperforms all other methods in a single step, which is a $100\times$ speedup compared to the diffusion models and a $1,024\times$ speedup compared to the language model in terms of number of function evaluations (NFE).

3 Enhancer DNA design

We now assess the performance of Dirichlet FM on DNA enhancer sequences and design evaluations that quantify both unconditional and conditional sample quality. Implementation and architecture details are in Appendix B.1.

Data. We evaluate on two enhancer sequence datasets from fly brain cells (Janssens et al., 2022) and from human melanoma cells (Atak et al., 2021). These contain 104k fly brain and 89k melanoma sequences of length 500 with cell class labels determined from ATAC-seq data (Buenrostro et al., 2013). Overall, there are 81 such classes of cells in the fly brain data and 47 in melanoma data (see Appendix B.2 for more data details).

Metric. To score the similarity between a data distribution and a generative model’s distribution, we employ a metric similar to the Fréchet inception distance (FID) that is commonly used to evaluate image generative models (Heusel et al., 2017) and was adapted to molecule generative models as Fréchet ChemNet distance (FCD) (Preuer et al., 2018). We follow this established principle and call our metric Fréchet Biological distance (FBD). Hence, we train a classifier model to predict cell types and use its hidden representations as embeddings of generated samples and data distribution samples. Then, the FBD is calculated as the Wasserstein distance between Gaussians fit to embeddings from the two distributions (10k each).

Q1: How well can Dirichlet FM capture the sequence distribution? We compare with an autoregressive language model (the best baseline in the promoter design experiments in Section 4.2) and with Linear FM. To evaluate, we calculate the FBD between the models’ generated sequences and the unconditional data distribution. Dirichlet FM outperforms the language model by a large margin on both datasets and linear FM for human melanoma cell enhancer generation (Table 2). Moreover, distillation minimally impacts FBD while speeding up inference by 3 orders of magnitude compared to the language model and 2 to Dirichlet FM (distilled Dirichlet FM only requires 1 step). Such speedups, compared to autoregressive models, make Dirichlet FM a promising direction for other applications with high sequence lengths where inference times are important.

Q2: Can guided Dirichlet FM produce class-specific sequences and improve upon the state-of-the-art? We condition Dirichlet FM on different target cell-type classes via guidance (Section 3.4). To quantify how well the generated sequences match the target class distribution, we use the FBD between the generated distribution and the data distribution conditioned on the target class. Additionally, we train a separate cell-type classifier and evaluate the probability it assigns to the target class for a generated sequence.

Sequences of classifier-free guided Dirichlet FM for the cell-type perineurial glia (PNG) have better FBD (Figure 5) than the sequences of Taskiran et al. (2023), of which several were experimentally validated as functioning enhancers (we show this comparison only for the PNG class since their sequences are available for it). For the other classes, guidance is similarly effective; by increasing the guidance factor $\gamma$, the classifier probability for generated sequences to belong to the target class can be improved close to 100%, and the FBD improves significantly until reaching a minimum (after the minimum the diversity decreases and the FBD worsens). The improvements with classifier guidance (Appendix Figure 9) are still significant but smaller.

Q3: Can guidance improve unconditional generation? We generate unconditional sequences with classifier-free guided Dirichlet FM by first sampling a class (based on its empirical frequency) and then guiding toward that class. This significantly improves sample quality compared to unguided Dirichlet FM (Figure 6) and baselines (Table 2). Thus, guided Dirichlet FM via the connection we derive between flow and score reproduces the success in image diffusion models of enhancing sample quality via conditioning (Rombach et al., 2022; Saharia et al., 2022). Furthermore, like guidance for images (Xu et al., 2023), the FBD first improves under increased guidance (Figure 6) until reaching a minimum, after which increased guidance deteriorates sample diversity and, therefore, FBD.

Conclusion

We presented Dirichlet flow matching for modeling discrete data via a generative process on the simplex. This solves the pathological behavior of linear flow matching on the simplex, which we identified. Compared with autoregressive methods or diffusion with discrete noise, Dirichlet FM enables distillation and conditional generation via guidance, which we derived via a connection between the marginal flow and the scores of a mixture of Dirichlets.

Experimental results on important regulatory DNA sequence design tasks across 3 datasets demonstrate Dirichlet FM’s effectiveness and utility for hard generative modeling tasks over long sequences. The results confirm Dirichlet FM’s superiority to linear FM and multiple discrete diffusion approaches. Distilling Dirichlet FM only marginally impacts performance while enabling one-step generation, leading to multiple orders of magnitude speedups compared to autoregressive generative models for long sequences. Lastly, we demonstrated effective class conditional generation via guided Dirichlet FM to design cell-type specific enhancers - an important task for gene therapies. Hence, Dirichlet FM is a flexible framework (guidance, distillation) with strong performance for biological sequence generation and a promising direction for further discrete data applications.

Acknowledgements

We thank Andrew Campbell, Yaron Lipman, Felix Faltings, Jason Yim, Ruochi Zhang, Rachel Wu, Jason Buenrostro, Bernardo Almeida, Gokcen Eraslan, Ibrahim I. Taskiran, and Pavel Avdeyev for helpful discussions.

This work was supported by the NSF Expeditions grant (award 1918839: Collaborative Research: Understanding the World Through Code), the Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium, the Abdul Latif Jameel Clinic for Machine Learning in Health, the DTRA Discovery of Medical Countermeasures Against New and Emerging (DOMANE) Threats program, the DARPA Accelerated Molecular Discovery program, the NSF AI Institute CCF-2112665, the NSF Award 2134795, and the GIST-MIT Research Collaboration grant.

References

Appendix A Method Details

For all $t$ and $\mathbf{x}\in S_{K}$, at convergence our denoising classifier satisfies

Thus, if we parameterize the vector field via Equation 10, then we are assured that

Proof of Proposition 1

Suppose that a flow matching model is trained with the linear flow map (Equation 11). Then, for all $k=2,\ldots K$ and $\mathbf{x}\sim p_{t}(\mathbf{x})$, the converged model posterior $p(\mathbf{x}_{1}\mid\mathbf{x})\propto p_{t}(\mathbf{x}\mid\mathbf{x}_{1})p_{\text{data}}(\mathbf{x}_{1})$ has support over at most $k-1$ vertices for times $t>1/k$.

In linear flow matching, the explicit form of the conditional probability path is given by

Suppose for sake of contradiction that $t>1/k$ but $p_{t}(\mathbf{x}\mid\mathbf{x}_{1})p_{\text{data}}(\mathbf{x}_{1})\neq 0$ for $k^{\prime}\geq k$ values of $\mathbf{x}_{1}$. Without loss of generality, suppose that $\mathbf{e}_{1},\ldots\mathbf{e}_{k}$ are $k$ of those values. Then, by Equation 27, we have $x_{i}\geq t$ for $i=1,\ldots k$. Then

This contradicts the fact that $\mathbf{x}$ must lie on the simplex. ∎

A.1 Dirichlet Conditional Vector Field

As preliminaries, we recall the definition of the multivariate beta function:

and that of the incomplete (two-argument) beta function:

with the identity $\mathcal{B}(a,b)=\mathcal{B}(1;a,b)$.

We wish to construct a conditional flow $u_{t}(\mathbf{x}\mid\mathbf{x}_{1}=\mathbf{e}_{i})$ which generates the evolution of the conditional probability path

We choose the following ansatz for the functional form of $u_{t}$:

i.e., (1) the flow points towards the target vertex $\mathbf{e}_{i}$ and (2) the magnitude is scaled by a constant dependent only on $x_{i}$ and $t$. Then consider the $(K-2)$-dimensional hyperplane $A_{b}$ of constant $x_{i}=b$ which cuts through the simplex, separating it into two regions $S_{1},S_{2}$ (Figure 7). We make the following key observation:

The probability flux crossing the plane $A_{b}$ is equal to the the rate of change of the total probability of region $S_{1}$.

Thus, we solve for the constant $C(x_{i},t)$ in the ansatz by deriving these two quantities and setting them equal to each other.

Q1: What is the probability mass of $S_{1}$ and its rate of change?

The probability mass of $S_{1}$ can be obtained by integrating the density over each hyperplane $A_{c}$ (defined by $x_{i}=c$) for $c<b$ and then integrating over all such hyperplanes. Since the density is a constant proportional to $x_{i}^{t}=c^{t}$ over each hyperplane, we obtain

where $\operatorname{Vol}(A_{c})$ refers to the volume of the intersection between the hyperplane and the simplex. Since this region is defined by $x_{i}=c$, the remaining entries of $\mathbf{x}$ must add up to $1-c$. Thus, $A_{c}$ can be viewed as a nonstandard probability simplex over $K-1$ variables. The volume of this $(K-2)$-dimensional space is proportional to $(1-c)^{K-2}$, giving

where $I$ is the so-called regularized incomplete beta function and is well-known as the CDF of the Beta distribution. Its derivative with respect to the first parameter is not available in closed form, but we write it as

Q2: What is the probability flux across the hyperplane $A_{b}$? The probability flux across $A_{b}$ (into $S_{2}$) is given by

The normal vector $\mathbf{n}$ points from the center of the face opposite $\mathbf{e}_{i}$, specified by $x_{i}=0$ and $x_{j}=1/(K-1),j\neq i$, towards $\mathbf{e}_{i}$. Thus, the probability flux density is given by the dot product of

Since $\sum_{j\neq i}x_{j}=1-b$, this dot product is equal to $(1-b)K/(K-1)$. We also see that $\lVert\mathbf{n}\rVert=\sqrt{K/(K-1)}$. Importantly, the flux density is constant on the hyperplane $A$, so the total flux is given by a simple product:

Checking the transport equation. We check that $\nabla\cdot(p_{t}u_{t})=-\partial p_{t}/\partial t$:

Appendix B Experimental details

We provide algorithms for training and inference in Algorithm 1, and Algorithm 2. These describe the procedures for generating discrete data on a single simplex. When generating a discrete data sequence, they are extended easily by performing the same procedure for a sequence of simplices. Code for reproducing all results is available in the accompanying supplementary material as .zip file.

Toy experiments. We train all models in Figure 4 for 450,000 steps with a batch size of 512 to ensure that they have all converged and then evaluate the KL of the final step.

Promoter Design. We follow the setup of Avdeyev et al. (2023) and train for 200 epochs with a learning rate of $5\times 10^{-4}$ and early stopping on the MSE on the validation set. We communicated with Avdeyev et al. (2023) to ensure that we have the same training and inference setup as them, and we build on their codebase to evaluate the generated sequences with the Sei regulatory activity prediction model (Chen et al., 2022). Thus, we use 100 inference steps for our Dirichlet FM instead of the 400 that they use in their code since they state that they used 100 integration steps for the results in the paper, which is also stated in the paper. Under this setup, we obtain the performance numbers for Linear FM and Dirichlet FM. We also ran the autoregressive language model in this setup (except for the number of inference steps which does not apply). Meanwhile, the numbers we report for Bit Diffusion, D3PM, and DDSM are taken from the DDSM paper (Avdeyev et al., 2023).

Enhancer Design. For both evaluations, on the human melanoma cell and the fly brain cell dataset, we train for 800 epochs (convergence of validation curves is reached after approximately 300 for both datasets). We use FBD for early stopping. For inference, we use 100 integration steps.

Classifier-free Guidance. For classifier-free guidance, we train with a conditioning ratio (the fraction of times we train with a class label as input instead of the no-class token as input) of 0.7. During inference, the inference Algorithm 2 is changed in that two probability vectors $\hat{\mathbf{p}}$ are predicted, once with class conditioning and once without class conditioning, which we sum together according to Equation 22. Then, we project the resulting probabilities onto the simplex since negative values can arise for guidance factors $\gamma>1$. For this purpose, we use the algorithm by Wang & Carreira-Perpinán (2013).

Classifier Guidance. The noisy classifier that we train for classifier guidance has the same architecture as our generative model, except that we sum the final representations and feed them into a 2-layer feed-forward network that serves as classification head. For training, we use early stopping on the accuracy and train for approximately 800 epochs. During inference, we use automatic differentiation to obtain the classifier’s gradients with respect to the input points on the simplices. To perform classifier guidance, we then convert the flow model output probabilities to scores as described in Section 3.4, obtain the guided score by adding the unconditional and the conditional scores via Equation 23, and convert the obtained scores back to probabilities. These we project to the simplex (Wang & Carreira-Perpinán, 2013) from which we obtain the vector field for integration.

Classifier for FBD calculation. This classifier’s architecture is similar to that of the noisy classifier for classifier guidance. However, it does not have any time conditioning and takes token embeddings as input instead of points on the simplex. For training, we use early stopping on the accuracy and train approximately 100 epochs.

The sequence embeddings that we use to calculate FBD are given by the hidden features after the first layer of the classification head. The 4 classes that we choose for the cell type specific enhancer generation are chosen as classes with a good tradeoff between the area under the receiver operator characteristic curve and the area under the precision-recall curve.

Distillation. For distillation, we run inference with the teacher model for every training step of the student model to obtain pairs of noise and training targets. With this, we train the student model on approximately 6 billion sequences.

Computational requirements. We train on RTX A600 GPUs. Training in the enhancer generation setup for 200 epochs on sequences with length 500 takes 7 hours. Our largest model uses approximately 8GB of RAM during training.

Architecture The architecture that we use for the promoter design experiments is the same as in DDSM (Avdeyev et al., 2023). In our other experiments we replaced group norm with layer norm in their architecture. The model consists of 20 layers of 1D convolutions interleaved with time embedding layers (and class-type embedding layers for classifier-free guidance) and normalization layers. We also experimented with Transformer architectures, which led to worse performance. All models use this 20-layer architecture except for the classifier for the Fly Brain data, which has 5 layers.

B.2 Data

For the enhancer data of 104665 fly brain cell sequences (Janssens et al., 2022), we use the same split as Taskiran et al. (2023), resulting in an 83726/10505/10434 split for train/val/test. Meanwhile, for the human melanoma cell dataset of 88870 sequences (Atak et al., 2021), their split has 70892/8966/9012 sequences. It is noteworthy that these datasets contain ATAC-seq data (Buenrostro et al., 2013), which means that not all sequences are guaranteed to be enhancers and actually enhance transcription of a certain gene. ATAC-seq only measures the chromatin accessibility of the sequences in the cell types, which is a necessary but not sufficient requirement for a sequence to be an enhancer. In Figure 8, we show histograms for the class distributions of both datasets. For the melanoma dataset, there is little class imbalance.

Appendix C Additional Results

Analytical toy experiment for classifier guidance. As a toy experiment to demonstrate our classifier guidance procedure, we construct a distribution conditioned on a binary random variable. The conditional distribution is a categorical distribution over 20 classes. In this setup, the time-dependent class probabilities conditioned on a noisy point on the simplex can be computed analytically. Thus, we can use them to obtain the vector field for Dirichlet FM analytically. Furthermore, the class probabilities and the gradients of their logarithm (the score) can be computed analytically. Hence, we can simulate classifier guided Dirichlet FM analytically for this toy distribution.

The results in Figure 10 show a close match between the empirical distribution of the generated data and the ground truth probabilities. This empirically confirms the effectiveness of our classifier guidance procedure that relies on converting probabilities to scores and converting them back to probabilities by solving a linear system of equations.