Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design

Introduction

Scientific domains often involve continuous atomic interactions with discrete chemical descriptions. Expanding the capabilities of generative models to handle discrete and continuous data, which we refer to as multimodal, is a fundamental problem to enable their widespread adoption in scientific applications (Wang et al., 2023). One such application requiring a multimodal generative model is protein co-design where the aim is to jointly generate continuous protein structures alongside corresponding discrete amino acid sequences (Shi et al., 2022). Proteins have been well-studied: the function of the protein is endowed through its structure while the sequence is the blueprint of how the structure is made. This interplay motivates jointly generating the structure and sequence rather than in isolation. To this end, the focus of our work is to develop a multimodal generative framework capable of co-design.

Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020) have achieved state-of-the-art performance across multiple applications. They have potential as a multimodal framework because they can be defined on both continuous and discrete spaces (Hoogeboom et al., 2021; Austin et al., 2021). However, their sample time inflexibility makes them unsuitable for multimodal problems. On even just a single modality, finding optimal sampling parameters requires extensive re-training and evaluations (Karras et al., 2022). This problem is exacerbated for multiple modalities. On the other hand, flow-based models (Liu et al., 2023; Albergo & Vanden-Eijnden, 2023; Lipman et al., 2023) improve over diffusion models with a simpler framework that allows for superior performance through sampling flexibility (Ma et al., 2024). Unfortunately, our current inability to define a flow-based model on discrete spaces holds us back from a multimodal flow model.

We address this by introducing a novel flow-based model for discrete data named Discrete Flow Models (DFMs) and thereby unlock a complete framework for flow-based multimodal generative modeling. Our key insight comes from seeing that a discrete flow-based model can be realized using Continuous Time Markov Chains (CTMCs). DFMs are a new discrete generative modeling paradigm: less restrictive than diffusion, allows for sampling flexibility without re-training and enables simple combination with continuous state space flows to form multimodal flow models.

Fig. 1A provides an overview of DFMs. We first define a probability flow $p_{t}$ that linearly interpolates from noise to data. We then generate new data by simulating a sequence trajectory $x_{t}$ that follows $p_{t}$ across time which requires training a denoising neural network with cross-entropy. The sequence trajectory could have many transitions or few, a property we term CTMC Stochasticity (Fig. 1B). Prior discrete diffusion models are equivalent to picking a specific stochasticity at training time, whereas we can adjust it at inference: enhancing sample quality and exerting control over sample distributional properties.

Using DFMs, we are then able to create a multimodal flow model by defining factorized flows for each data modality. We apply this capability to the task of protein co-design by developing a novel continuous structure and discrete sequence generative model named Multiflow. We combine a DFM for sequence generation and a flow-based structure generation method developed in Yim et al. (2023a). Previous multimodal approaches either generated only the sequence or only the structure and then used a prediction model to infer the remaining modality (see Sec. 5). Our single model can jointly generate sequence and structure while being able to condition on either modality.

In our experiments (Sec. 6), we first verify on small scale text data that DFMs outperform the discrete diffusion alternative, D3PM (Austin et al., 2021) through their expanded sample time flexibility. We then move to our main focus, assessing Multiflow’s performance on the co-design task of jointly generating protein structure and sequence. Multiflow achieves state-of-the-art co-design performance while data distillation allows for obtaining state-of-the-art structure generation. We find CTMC stochasticity enables controlling sample properties such as secondary structure composition and diversity. Preliminary results on inverse and forward folding show Multiflow is a promising path towards a general-purpose protein generative model.

Our contributions are summarized as follows:

We present Discrete Flow Models (DFMs), a novel discrete generative modeling method built through a CTMC simulating a probability flow.

We combine DFMs with continuous flow-based methods to create a multimodal generative modeling framework.

We use our multimodal framework to develop Multiflow, a state-of-the-art generative protein co-design model with the flexibility of multimodal protein generation.

Background

We aim to model discrete data where a sequence $x\in\{1,\dots,S\}^{D}$ has $D$ dimensions, each taking on one of $S$ states. For ease of exposition, we will assume $D=1$; all results hold for $D>1$ as discussed in App. E. We first explain a class of continuous time discrete stochastic processes called Continuous Time Markov Chains (CTMCs) Norris (1998) and then describe the link to probability flows.

In practice, we need to simulate the sequence trajectory with finite time intervals $\Delta t$. A sequence trajectory can be simulated with Euler steps (Sun et al., 2023b)

where the sequence starts from an initial sample $x_{0}\sim p_{0}$ at time $t=0$. The rate matrix $R_{t}$ along with an initial distribution $p_{0}$ together define the CTMC.

2 Kolmogorov equation

For a sequence trajectory following the dynamics of a CTMC, we write its marginal distribution at time $t$ as $p_{t}(x_{t})$. The Kolmogorov equation allows us to relate the rate matrix $R_{t}$ to the change in $p_{t}(x_{t})$. It has the form:

The difference between the incoming and outgoing probability mass is the time derivative of the marginal $\partial_{t}p_{t}(x_{t})$. Using our definition of $R_{t}(x_{t},x_{t})$, Eq. 5 can be succinctly written as $\partial_{t}p_{t}=R_{t}^{\top}p_{t}$ where the marginals are treated as probability mass vectors: $p_{t}\in^{S}$. This defines an Ordinary Differential Equation (ODE) in a vector space. We refer to the series of distributions $p_{t}\ \forall t\in$ satisfying the ODE as a probability flow.

Discrete Flow Models

A Discrete Flow Model (DFM) is a Discrete data generative model built around a probability Flow that interpolates from noise to data. To sample new datapoints, we simulate a sequence trajectory that matches the noise to data probability flow. The flow construction allows us to combine DFM with continuous data flow models to define a multimodal generative model. Proofs for all propositions are in App. B.

The conditional flow, $p_{t|1}(\cdot|x_{1})$ interpolates from noise to the datapoint $x_{1}$. The conditioning allows us to write the flow down in closed form. We are free to define $p_{t|1}(\cdot|x_{1})$ as needed for the specific application. The conditional flows we use in this paper linearly interpolate towards $x_{1}$ from a uniform prior or an artificially introduced mask state, $M$:

If $R_{t}(x_{t},j|x_{1})$ is a rate matrix that generates the conditional flow $p_{t|1}(x_{t}|x_{1})$, then

Our aim now is to calculate $R_{t}(x_{t},j|x_{1})$ and $p_{1|t}(x_{1}|x_{t})$ to plug into Eq. 9. $p_{1|t}(x_{1}|x_{t})$ is the distribution predicting clean data $x_{1}$ from noisy data $x_{t}$ and in Sec. 3.1.2, we will train a neural network $p_{1|t}^{\theta}(x_{1}|x_{t})$ to approximate it. In Sec. 3.2, we will show how to derive $R_{t}(x_{t},j|x_{1})$ in closed form. Sampling pseudo-code is provided in Alg. 1.

We discuss further CTMC sampling methods in App. G. Our construction of the generative flow from conditional flows is analogous to the construction of generative probability paths from conditional probability paths in Lipman et al. (2023), where instead of a continuous vector field generating the probability path, we have a rate matrix generating the probability flow. We expand on these links in Table. 1.

1.2 Training

We train a neural network with parameters $\theta$, $p_{1|t}^{\theta}(x_{t}|x_{1})$, to approximate the true denoising distribution using the standard cross-entropy i.e. learning to predict the clean datapoint $x_{1}$ when given noisy data $x_{t}\sim p_{t|1}(x_{t}|x_{1})$.

2 Choice of Rate Matrix

The missing piece in Eq. 9 is a conditional rate matrix $R_{t}(x_{t},j|x_{1})$ that generates the conditional flow $p_{t|1}(x_{t}|x_{1})$. There are many choices for $R_{t}(x_{t},j|x_{1})$ that all generate the same $p_{t|1}(x_{t}|x_{1})$ as we later show in Prop. 3.3. In order to proceed, we start by giving one valid choice of rate matrix and from this, build a set of rate matrices that all generate $p_{t|1}$. At inference time, we can then pick the rate matrix from this set that performs the best. Our starting choice for a rate matrix that generates $p_{t|1}$ is defined for $x_{t}\neq j$ as,

Assuming zero mass states, $p_{t|1}(j|x_{1})=0$, have $\partial_{t}p_{t|1}(j|x_{1})=0$, then $R^{*}_{t}$ generates $p_{t|1}(x_{t}|x_{1})$.

as we derive in App. F. Using $R^{*}_{t}$ as a starting point, we now build out a set of rate matrices that all generate $p_{t|1}$. We can accomplish this by adding on a second rate matrix that is in detailed balance with $p_{t|1}$.

3 DFMs Recipe

We now summarize the key steps of a DFM. PyTorch code for a minimal DFM implementaton is provided in App. F.

Define the desired noise schedule $p_{t|1}(x_{t}|x_{1})$ (Sec. 3.1).

Train denoising model $p_{1|t}^{\theta}(x_{1}|x_{t})$ (Sec. 3.1.2).

Choose rate matrix $R_{t}^{\eta}$ (Sec. 3.2).

Multimodal Protein Generative Model

We now use a DFM to create a multimodal protein generative model. To generate multimodal data, we will define a multimodal generative flow. We define $p_{t|1}$ to factorize over different modalities allowing us to define $p_{t|1}$ individually for each one. Our training loss is then simply the sum of the standard flow loss for each modality. At inference time, we can also update each modality individually for each simulation step, using an ODE for continuous data and a CTMC for discrete data. We now apply this capability on protein structure-sequence generation.

A protein can be modeled as a linear chain of residues, each with an assigned amino acid and 3D atomic coordinates. Protein co-design aims to jointly generate the amino acids (sequence) and coordinates (structure). Prior works have used a generative model on one modality (sequence or structure) with a separate model to predict the other (see Sec. 5). Instead, our approach uses a single generative model to jointly sample both modalities: a DFM for the sequence and a flow model, FrameFlow (Yim et al., 2023a), for the structure. We refer to this as co-generating the sequence and structure; hence, the method is called Multiflow.

Sampling. To sample with Multiflow, we integrate along the ODE trajectories for the translations and rotations whilst simultaneously following the CTMC for the amino acid sequence. Each Euler step during sampling has the update:

Related Work

Discrete Diffusion Models. Our continuous time flow builds on work that extends discrete diffusion Hoogeboom et al. (2021); Austin et al. (2021) to continuous time Campbell et al. (2022); Sun et al. (2023b); Santos et al. (2023); Lou et al. (2023) but we simplify and extend the framework. We are not restricted to noising processes that can be defined by a matrix exponential as we just write $p_{t|1}$ down directly and we have the freedom to choose $R_{t}(x_{t},j|x_{1})$ at inference time rather than being restricted to the time reversal. We show how DFMs encompasses prior discrete diffusion models in App. H. For molecular retrosynthesis, Igashov et al. (2023) also considered a data conditional process, but did not build a modeling framework around it. Zhang et al. (2023) constructed low-stochasticity rate matrices and their derivation provides the building blocks of Prop. 3.2. Some works have built a multimodal diffusion model for molecule generation (Peng et al., 2023; Vignac et al., 2023b; Hua et al., 2023) whereas we focus on protein co-design using flows. We discuss further related work in App. D.

Protein Generation. Diffusion and flow models have risen in popularity for generating novel and diverse protein backbones (Yim et al., 2023b; a; Bose et al., 2023; Lin & AlQuraishi, 2023; Ingraham et al., 2023). RFDiffusion achieved notable success by generating proteins validated in wet-lab experiments (Watson et al., 2023). However, these methods required a separate model for sequence generation. Some works have focused only on sequence generation with diffusion models (Alamdari et al., 2023; Gruver et al., 2023; Yang et al., 2023; Yi et al., 2023). We focus on co-design which aims to jointly generate the structure and sequence.

Prior works have attempted co-design. ProteinGenerator (Lisanza et al., 2023) performs Euclidean diffusion over one-hot amino acids while predicting the structure at each step with RosettaFold (Baek et al., 2021). Conversely, Protpardelle (Chu et al., 2023) performs Euclidean diffusion over structure while iteratively predicting the sequence. Multiflow instead uses a generative model over both the structure and sequence which allows for flexibility in conditioning at inference time (see Sec. 6.2.1). Luo et al. (2022); Shi et al. (2022) are co-design methods, but are limited to generating CDR loops on antibodies. Lastly, Anand & Achim (2022) presented diffusion on structure and sequence, but did not report standard evaluation metrics nor is code available.

Experiments

We first show that tuning stochasticity at sample time improves pure discrete generative modeling performance by modeling text data. We then evaluate Multiflow, the first flow model on discrete and continuous state spaces. We show Multiflow provides state-of-the-art-performance on protein generation compared to prior approaches that do not generate using a true multimodal generative model. Finally, we investigate Multiflow’s crossmodal properties of how varying the sequence sampling affects the structure.

Results. Text samples are evaluated following Strudel et al. (2022). A much larger text model, we use GPT-J-6B Wang & Komatsuzaki (2021), is used to evaluate the negative log-likelihood (NLL) of the generated samples. The NLL metric alone can be gamed by repeating similar sequences, so the token distribution entropy is also measured. Good samples should have both low NLL and entropy close to the data distribution. For a given value of $\eta$, we create a Pareto-frontier in NLL vs entropy space by varying the temperature applied to the $p_{1|t}^{\theta}(x_{1}|x_{t})$ logits during the softmax operation. Fig. 2 plots the results for varying levels of $\eta$ and sampling temperature. For comparison, we also include results for the discrete diffusion D3PM method with absorbing state corruption Austin et al. (2021). We find the DFM performs better than D3PM due to our additional sample time flexibility. We are able to choose the value of $\eta$ that optimizes the Pareto-frontier at sample time (here $\eta=15$) whereas D3PM does not have this flexibility. We show the full $\eta$ sweep in App. I and show the frontier for $\eta=0$ in Fig. 2. When $\eta=0$, performance is similar due to DFMs being a continuous time generalization of D3PM at this setting, see Sec. H.2. We also include results for an autoregressive model in Fig. 2 for reference; however, we note this is not a complete like-for-like comparison as autoregressive models require much less compute to train than diffusion based models Gulrajani & Hashimoto (2023).

2 Protein generation

Metrics. Evaluating the quality of structure-sequence samples is performed with self-consistency which measures how consistent a generated sequence is with a generated structure by testing how accurately a protein folding network can predict the structure from the sequence. Specifically, either AlphaFold2 (Jumper et al., 2021) or ESMFold (Lin et al., 2023), is first used to predict a structure given only the generated sequence. Our results will use ESMFold but we show results with AlphaFold2 in App. J. Then, we calculate scRMSD: the Root Mean Squared Deviation between the generated and predicted structure’s backbone atoms. The generated structure is called designable if $\text{scRMSD}<2\text{\AA{}}$.

Structure-only generative models such as RFdiffusion first use ProteinMPNN (PMPNN) (Dauparas et al., 2022) to predict a sequence given the generated structure in order to then be able to use the self-consistency metric. We present three variants of self-consistency:

Co-design 1: use the sampled (structure, sequence) pair.

PMPNN 8: take only the sampled structure and predict 8 sequences with PMPNN. Then use ESMFold to predict a new structure for each sequence. The final structure-sequence pair is the original sampled structure along with the PMPNN sequence with minimum scRMSD.

PMPNN 1: same as PMPNN 8 except PMPNN only generates one sequence.

PMPNN 8 and PMPNN 1 evaluate only the quality of a model’s generated structures whereas, for co-design models, Co-design 1 evaluates the quality of a model’s generated (structure, sequence) pairs. The comparison between PMPNN 1 and Co-design 1 allows for evaluating the quality of co-designed sequences. PMPNN 8 is the procedure used in prior structure-only works. As our main metric of sample quality, we report designability as the percentage of designable samples. As a further sanity check, designable samples are then evaluated on diversity and novelty. We use FoldSeek (van Kempen et al., 2022) to report diversity as the number of unique clusters while novelty is the average TM-score (Zhang & Skolnick, 2005) of each sample to its most similar protein in PDB.

Training. Our training data consisted of length 60-384 proteins from the Protein Data Bank (PDB) (Berman et al., 2000) that were curated in Yim et al. (2023b) for a total of 18684 proteins. Training took 200 epochs over 3 days on 4 A6000 Nvidia GPUs using the AdamW optimizer (Loshchilov & Hutter, 2017) with learning rate 0.0001.

Distillation. Multiflow with PDB training generated highly designable structures. However, the co-designed sequences suffered from lower designability than PMPNN. Our analysis revealed the original PDB sequences achieved worse designability than PMPNN. We sought to improve performance by distilling knowledge from other models. To accomplish this, we first replaced the original sequence of each structure in the training dataset with the lowest scRMSD sequence out of 8 generated by PMPNN conditioned on the structure. Second, we generated synthetic structures of random lengths between 60-384 using an initial Multiflow model and added those that passed PMPNN 8 designability into the training dataset with the lowest scRMSD PMPNN sequence. We found that we needed to add only an extra 4179 examples to the original set of 18684 proteins to see a dramatic improvement. This procedure can be seen as a single step of reinforced self training (ReST) Gulcehre et al. (2023).

Following RFdiffusion’s benchmark, we sample 100 proteins for each length 70, 100, 200, and 300. We sample Multiflow with 500 timesteps using a temperature of 0.1 (PMPNN also uses 0.1) and stochasticity level $\eta=20$. We compare our structure quality to state-of-the-art structure generation method RFdiffusion. For co-design, we compare to Protpardelle and ProteinGenerator. All methods were ran using their publicly released code and evaluated identically.

Our results are presented in Table. 3. We find that Multiflow’s co-design capabilities surpass previous co-design methods, none of which use a joint multimodal generation process. Multiflow generates sequences that are consistent with the generated structure at a comparable level to PMPNN which we see through comparing the Co-design 1 and PMPNN 1 designability. On pure structure generation, we find that Multiflow outperforms all baselines in terms of structure quality measured by PMPNN 8 designability. Multiflow also attains comparable diversity and novelty to previous approaches. We ablate our use of distillation and find that distillation results in overall designability improvements while also improving diversity. Finally, we train our exact same architecture except only modeling the structure on the distilled dataset using the loss presented in Yim et al. (2023a). We find our joint structure-sequence model achieves the same structural quality as the structure-only version, however, additionally including the sequence in our generative process induces extra structural diversity.

Crossmodal modulation. We next investigate how modulating the CTMC stochasticity of the sequence affects the structural properties of sampled proteins. Fig. 3 shows that varying the stochasticity level $\eta$ results in a change of the secondary structure composition (Kabsch & Sander, 1983) of the sampled proteins. This is an example of the flexibility our multimodal framework provides to tune properties between data modalities at inference time.

2.2 Forward and Inverse Folding

Multiflow can achieve state-of-the-art codesign performance, but can accomplish more tasks as described in Fig. 1B and Table. 4. Expanding Multiflow to achieve competitive performance on all tasks is a future work. Here, we take the same model weights for co-design and evaluate forward and inverse folding without additional training. We compare performance to ESMFold and ProteinPMNN which are specialized models for forward and inverse folding. We curated a clustered test-out set of 449 monomeric proteins with length $<400$ from the PDB using a date split of our training set. Details of forward/inverse folding and these experiments can be found in App. J. We find Multiflow can achieve very close performance with ProteinMPNN while it achieves poor results compared to ESMFold. This highlights a limitation that Multiflow cannot perform competitively at every generation task, but leaves exciting future work for a potential general-purpose generative model.

Discussion

We presented Discrete Flow Models (DFMs), a flow based generative model framework by making analogy to continuous state space flow models. Our formulation is simple to implement, removes limitations in defining corruption processes, and provides more sampling flexibility for improved performance compared to previous discrete diffusion models. Our framework enables easy application to multimodal generative problems which we apply to protein co-design. The combination of a DFM and FrameFlow enables state-of-the-art co-design with Multiflow. Future work includes to develop more domain specific models with DFMs and improve Multiflow’s performance on all protein generation tasks including sidechain modeling.

Acknowledgments

The authors would like to thank Ricardo Baptista, Mathieu Le Provost, George Deligiannidis, Joe Benton, Bowen Jing, Hannes Stärk, Emile Mathieu, Luhuan Wu, Timur Garipov, Rachel Wu, Mingyu Choi, Sidney Lisanza, and Woody Ahern for helpful discussions.

AC acknowledges support from the EPSRC CDT in Modern Statistics and Statistical Machine Learning (EP/S023151/1) JY was supported in part by an NSF-GRFP. JY, RB, and TJ acknowledge support from NSF Expeditions grant (award 1918839: Collaborative Research: Understanding the World Through Code), 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 and the Sanofi Computational Antibody Design grant. IF is supported by the Office of Naval Research, the Howard Hughes Medical Institute (HHMI), and NIH (NIMH-MH129046). The authors would like to acknowledge the use of the University of Oxford Advanced Research Computing (ARC) facility in carrying out this work. http://dx.doi.org/10.5281/zenodo.22558.

Impact statement

In this paper we work to advance general purpose generative modeling techniques, specifically those used for modeling discrete and multimodal data. We apply these techniques to the task of protein generation. Improving protein modeling capabilities can have wide ranging societal impacts and care must be taken to ensure these impacts are positive. For example, improved modeling capabilities can help design better enzymes and drug candidates that can then go on to improve the lives of many people. Conversely, these general purpose techniques could also be misused to design toxic substances. To mitigate these risks, we do not present any specific methods to apply Multiflow to tasks that could be easily adjusted to the design of harmful substances without expert knowledge.

References

Appendix A Organization of Appendix

The Appendix is organized as follows. App. B provides proofs for all propositions in the main text. App. C analyses the cross entropy objective used to train DFM and links controlling the cross entropy to controlling the model log-likelihood. App. D discusses further related work. App. E shows how DFM can be applied to multidimensional data through applying factorization assumptions to $p_{t|1}$. App. F gives concrete realizations with PyTorch code for DFM using the masking or uniform forms for $p_{t|1}$. App. G discusses methods for sampling from CTMCs and discusses their relation to our sampling method. App. H compares DFM to classical discrete diffusion models in discrete and continuous time finding that they can be fit within the DFM framework. App. I gives further details and results for our text experiment. App. J gives further details and results for our protein co-design experiments.

Appendix B Proofs

When writing rate matrices, $R_{t}(i,j)$, we will assume $i\neq j$ unless otherwise explicitly stated.

We write $R_{t}(i)\vcentcolon=\sum_{j\neq i}R_{t}(i,j)$.

B.1 Proof of Proposition 3.1

B.2 Proof of Proposition 3.2

In the main text we provided the form for $R^{*}_{t}$ under the assumption that $p_{t|1}(j|x_{1})>0$ for all $j$. Before proving Prop. 3.2, we first give the full form for $R^{*}_{t}$. First, assuming $x_{t}\neq j$ and $p_{t|1}(x_{t}|x_{1})>0$ we have,

For our proof, we assume that $p_{t|1}(j|x_{1})=0\implies\partial_{t}p_{t|1}(j|x_{1})=0$. This assumption means that when we have dead states with zero probability mass, they cannot be resurrected and gain probability mass in the future. We begin the proof with the Kolmogorov equation for processes conditioned on $x_{1}$,

We will now verify that $R^{*}_{t}$ satisfies this Kolmogorov equation and thus generates the desired $p_{t|1}(x_{t}|x_{1})$ conditional flow. We will first check that the Kolmogorov equation is satisfied when $p_{t|1}(x_{t}|x_{1})>0$. With this form of rate matrix, the RHS of equation (30) becomes

In the case that $p_{t|1}(x_{t}|x_{1})=0$ by assumption we have that $\partial_{t}p_{t|1}(x_{t}|x_{1})=0$. We have both $R^{*}_{t}(x_{t},j|x_{1})=0$ and $R^{*}_{t}(j,x_{t}|x_{1})=0$ because $p_{t|1}(x_{t}|x_{1})=0$. Therefore we have $\text{LHS}=\text{RHS}=0$ and thus the Kolmogorov equation is satisfied.

Intuitively, we require the assumption that dead states cannot be resurrected because $R^{*}_{t}$ is designed such that all states can equally distribute the mass flux requirements of making sure the marginal derivatives $\partial_{t}p_{t|1}(x_{t}|x_{1})$ are satisfied. If there is a state for which $p_{t|1}(x_{t}|x_{1})=0$ but $\partial_{t}p_{t|1}(x_{t}|x_{1})>0$ then this state would require mass from other states but could not provide any mass of its own since $p_{t|1}(x_{t}|x_{1})=0$. This would then violate the sharing symmetry required for our form of $R^{*}_{t}$. We note that this assumption is not strictly satisfied for the masking interpolant at $t=0$ or $t=1$ and not satisfied for the uniform interpolant at $t=1$. However, it is satisfied for any $t\in(0,1)$ and so we can conceptualize starting our process at $t=\epsilon$, $\epsilon\ll 1$, $\epsilon>0$, approximating a sample from $p_{\epsilon}(x_{\epsilon})$ with a sample from $p_{0}(x_{0})$ and running the process until $t=1-\epsilon$ and stopping here. The approximation can be made arbitrarily accurate by taking $\epsilon\rightarrow 0$.

B.3 Proof of Proposition 3.3

A rate matrix that satisfies the detailed balance condition (13) will result in $\partial_{t}p_{t|1}(i|x_{1})=0$ when simulating with this rate. This can be seen by substituting into the conditional Kolmogorov equation (30)

B.4 Proof of Proposition 3.4

We will assume we have $D$ dimensional data $x_{1}^{1:D}$ with each $x_{1}^{d}\in\{1,\dots,S\}$. We give an overview of how our method operates in the multi-dimensional case in Appendix E. Namely, we assume that our conditional flow factorizes as $p_{t|1}(x_{t}^{1:D}|x_{1}^{1:D})=\prod_{d=1}^{D}p_{t|1}(x_{t}^{d}|x_{1}^{d})$. We also assume that our rate matrix is for jumps that vary more than $1$ dimension at a time. Our optimality results are derived under these assumptions.

We first prove that $R^{*}_{t}$ achieves the minimum number of transitions for the masking interpolant case. We have

with $\mathcal{Z}_{t}^{d}=|\{j^{d}:p_{t|1}(j^{d}|x_{1}^{d})>0\}|$. Substituting in $\partial_{t}p_{t|1}$ and $p_{t|1}$ in the masking case gives

We refer to Appendix F.1 for the details of this derivation. Since ${R^{*}_{t}}^{d}$ depends only on $x_{t}^{d}$, $j^{d}$ and $x_{1}^{d}$ and not values in any other dimensions, each dimension propagates independently and we can consider each dimension in isolation. Consider the process for dimension $d$. The CTMC begins in state $x_{0}^{d}=M$. We have ${R^{*}_{t}}^{d}(x_{t}^{d}=M,j^{d}|x_{1}^{d})=\frac{1}{1-t}\delta\left\{j^{d},x_{1}^{d}\right\}$. Therefore, the only possible next state that the process can jump to is $x_{1}^{d}$. Once the process has jumped to $x_{1}^{d}$, the rate then becomes ${R^{*}_{t}}^{d}(x_{t}^{d}=x_{1}^{d},j^{d}|x_{1}^{d})=0$. We also know that the process must jump because $p_{1}(x_{t}^{d}|x_{1}^{d})=\delta\left\{x_{t}^{d},x_{1}^{d}\right\}$, $x_{1}^{d}\neq M$ and we know our rate matrix traverses our desired marginals by Proposition 3.2. Therefore, exactly one jump is made in dimension $d$. In total, our $D$ dimensional process will make $D$ jumps. Under our factorization assumption, during a jump no more than one dimension can change value. Therefore, the absolute minimum number of jumps for any process that starts at $x_{0}^{1:D}$ with $x_{0}^{d}=M,\forall d$ and ends at $x_{1}^{1:D}$, $x_{1}^{d}\neq M,\forall d$ is $D$. Our prior distribution is $p_{0}(x_{0}^{d})=\delta\left\{x_{0}^{d},M\right\}$ and so for any $x_{0}$ sample, we will always need to make $D$ jumps. Therefore, the minimum expected number of jumps is $D$ and $R^{*}_{t}$ achieves this minimum.

B.4.2 Uniform Interpolant

We now prove that $R^{*}_{t}$ achieves the minimum number of transitions for the uniform interpolant case. The conditional flow is

With this interpolant, our rate matrix becomes

B.4.3 Discussion

We have proven $x_{1}$ conditioned optimality only for the two simple conditional flows featured in the main text and we note that this result in not generally true for any conditional flow. Intuitively this is because $R^{*}_{t}$ treats the distribution of mass symmetrically between states, considering only the local differences in $\partial_{t}p_{t|1}$ between pairs of states. In general, the optimal rate would need to solve a global programming problem.

Appendix C Analysis of Training Objective

Before beginning the proof, we introduce path space measures for CTMC processes, following the exposition in Del Moral & Penev (2017), Chapter 18. A path of a CTMC is a single trajectory from time to time $t$. The trajectory is a function $\omega:s\in[0,t]\mapsto\omega_{s}\in\{1,\dots,S\}$ that is everywhere right continuous and has left limits everywhere (also known as càdlàg paths). Intuitively, it is a function that takes in a time variable and outputs the position of the particle following the trajectory at that time. The càdlàg condition in our case states that at jump time $\tau$ we have $\omega_{\tau}$ taking the new jumped to value and $\omega_{\tau}^{-}\vcentcolon=\lim_{s\uparrow\tau}\omega_{s}$ being the previous value before the jump, see Fig. 1B.

For an exponential random variable with time-inhomogeneous rate, the cumulative distribution function is given by

In other words, we simply start a new exponential timer once the previous jump occurs and the same equation carries through.

where $p_{0}$ is the initial state distribution.

We begin by marginalizing out the latent variables, $W_{t}$, $t\in[0,1)$ for our generative CTMC

We now apply Girsnov’s transformation using our $x_{1}$ conditioned CTMC

we note at this point that $p_{0|1}(W_{0}|x_{1})=p_{0}(W_{0})$ and the two intial distribution terms cancel out. Now, apply Jensen’s inequality

and take the expectation with respect to the data distribution

where we have used our definition of the jump distribution of

Substituting this into our form for $\mathcal{L}_{\text{ELBO}}$ given in equation (79) gives

Substituting this into our original bound on the model log-likelihood gives

and $C$ is a constant term independent of $\theta$.

In the next stages of the proof, we going to show that $\mathcal{L}_{\text{ce}}$ is the weighted cross-entropy, $\mathcal{L}_{R}$ is a regularizer towards the arbitrarily chosen $x_{1}$ conditioned rate matrix that we argue we can ignore and $\mathcal{L}_{\text{KL}}$ is a KL term that we will absorb into the bound on the model log-likelihood.

In order to proceed, we will need to make use of Dynkin’s formula

where $f(\cdot,\cdot)$ is a two-argument function. This formula can be understood intuitively as allowing us to switch from a sum over the jump times to a full integral over the time interval appropriately weighted by the probability that a jump occurs and the destination to which a jump goes to.

We first show that $\mathcal{L}_{\text{ce}}$ is the weighted cross entropy.

We now analyse the term $\mathcal{L}_{R}$. We will show that it is approximately equal to an objective which at its optimum sets the learned generative rate matrix to have the same overall jump probability as the arbitrarily chosen rate matrix that generates our $p_{t|1}(x_{t}|x_{1})$ conditional flow.

We will denote the approximate form of $\mathcal{L}_{R}$ as $\hat{\mathcal{L}}_{R}$.

We then split $\mathcal{L}_{\text{ELBO}}$ into three terms $\mathcal{L}_{\text{ce}}+\mathcal{L}_{R}+\mathcal{L}_{\text{KL}}$. We have seen how the term $\mathcal{L}_{\text{KL}}$ allows us to remove the standard $\mathcal{L}_{\text{ELBO}}$ training signal for the denoising model $p_{\theta}(x_{1}|x_{t})$ and replace it with the cross entropy, creating the $\mathcal{L}_{\text{ce}}$ term. This creates a looser bound if we are to train without the $\mathcal{L}_{\text{KL}}$ term,

We then argue that $\mathcal{L}_{R}$ is close to $\hat{\mathcal{L}}_{R}$ which is an unnecessary forcing term encouraging our generative rate to achieve a similar jump rate to our chosen $R_{t}(x_{t},j|x_{1})$ even though this $R_{t}$ matrix is an arbitrary decision and will have a different jump rate depending on which $R_{t}$ is chosen. We are then left with the standard cross entropy term as our final objective for $p_{\theta}(x_{1}|x_{t})$ with a final modification to its unweighted form for implementation ease.

C.2.1 Objective for the Masking Interpolant

In this section we will show that $\mathcal{L}_{\text{ELBO}}$ is exactly the weighted cross entropy for the case when we use the masking form for $p_{t|1}(x_{t}|x_{1})$. We note that a similar result has been proven by Austin et al. (2021) for the discrete time diffusion model, and here we verify that this result also holds for our DFM model. We will assume multi-dimensional data, $x_{1}\in\{1,\dots,S\}^{D}$. We refer to Appendix E for the details of the multi-dimensional setting. We will also assume that we use $R^{*}_{t}$ as our rate matrix that generates the $p_{t|1}(x_{t}|x_{1})$ conditional flow.

Before we manipulate $\mathcal{L}_{\text{ELBO}}$, we will first find the forms of $R^{*}_{t}(x_{t}^{1:D},j^{1:D}|x_{1}^{1:D})$, $R_{t}^{\theta}(x_{t}^{1:D},j^{1:D})$ and $R_{t}^{\theta}(x_{t}^{1:D})$ for the masking case. From Appendix F.1, equation (209) we have,

From Appendix F.1, equation (212) we have that,

We now find $R_{t}^{\theta}(x_{t}^{1:D})$

where on the final line we have used the fact that $p_{\theta}(x_{1}^{d}=M|x_{t}^{1:D})=0$.

We are now ready to manipulate the form of $\mathcal{L}_{\text{ELBO}}$. We start with

We now substitute in the masking forms for $R_{t}^{\theta}(W_{t})$, $R^{*}_{t}(W_{t},y|x_{1})$ and $R_{t}^{\theta}(W_{t},y)$

where we have moved terms that don’t depend on $\theta$ into the constant.

where we have arrived at the weighted cross entropy, weighted by $\frac{1}{1-t}$ and only calculated for dimensions that are masked in our corrupted sample $x_{t}$.

Appendix D Discussion of Related Work

There have been post-hoc changes to the sampling process made in prior work e.g. corrector steps used by Campbell et al. (2022), however due to the ELBO maximizing the model log-likelihood under the assumption of sampling using the time-reversal, the diffusion framework still revolves around one ‘canonical’ sample time process (the time-reversal) whereas DFM makes it clear this choice is arbitrary and the sample process can be chosen at inference time for best performance.

Previous discrete diffusion works have also suggested alternatives to the ELBO. Sun et al. (2023b) introduce a categorical score matching loss that resembles the cross entropy, however, the denoising network is required to make a prediction $x_{0}^{d}$ based only on the other $D-1$ dimensions of the input noisy state, $x_{t}^{1:D\backslash d}$. This requires specialized architectures and methods to remain computationally efficient. Vignac et al. (2023a) propose to learn a diffusion based model solely using the cross-entropy but do not analyse the link between the cross-entropy and the log-likelihood of the model as we do in App. C. Meng et al. (2022) propose to learn a discrete score model based on data ratios using an L2 based loss which has some undesirable properties such as not penalizing mode dropping as described by Lou et al. (2023). Lou et al. (2023) refine this approach and propose to learn data ratios using the score entropy loss which, like the standard cross entropy, does not depend on the choice of forward rate matrix. However, in order for the score entropy to be a true ELBO, the forward rate matrix needs to be used as a weighting factor.

Multimodal diffusion models have been applied to tabular data Kotelnikov et al. (2023) where continuous diffusion is used for continuous features and a uniform style of corruption under a discrete diffusion framework is applied to discrete features. This idea was then expanded to molecule generation where the task is to generate a molecules atom types, their positions and their connectivity. Peng et al. (2023) use a masking process for the discrete atom types and bond types with a continuous space process for the atom positions. Vignac et al. (2023b) use a discrete process converging towards the independent marginal distribution in each dimension (Vignac et al., 2023a) for atom types, bond types and formal charges of the molecules along with a continuous process for atom positions. Hua et al. (2023) use a uniform discrete process for bond types with a continuous space process applied to atom positions as well as atom features embedded in continuous space. These works also investigate the importance of the multimodal noise schedule. Peng et al. (2023) find that corrupting the bonds first and then the atom positions improves performance by avoiding unphysical bonds appearing in the corruption process. Vignac et al. (2023b) have a similar finding that during corruption, the atom types should be corrupted first, then the bond types and finally the atom positions. We generalize these ideas by using the approach of Albergo et al. (2023) and learning our model over all relative levels of noise between our modalities. This allows picking the desired path through the multimodal noise landscape at inference time either performing co-generation, inverse folding or forward folding.

Other approaches for discrete data modelling opt to embed the data into a continuous space in order to still use the continuous diffusion framework Li et al. (2022); Chen et al. (2023); Richemond et al. (2022); Gong et al. (2023); Dieleman et al. (2022); Han et al. (2022); Strudel et al. (2022); Gulrajani & Hashimoto (2023); Floto et al. (2023), however, this loses the discrete structure of the data during generation. This can be important when the quantity that is represented by the discrete variable as algorithmic importance. For example, Qin et al. (2023) perform sparse graph generation where the discrete token represents the existence of an edge. It is then important for the edge to be known to physically exist or not so that sparse graph networks can be applied to the problem.

General Fokker-Planck equations on discrete state spaces Chow et al. (2012) have been used to construct sampling methods for energy functions Sun et al. (2023a). Further, in a generative modelling context, the Kolmogorov equation has been used to construct equivalent diffusion processes with fewer transitions Zhang et al. (2023) making links to optimal transport. We take this idea further to build a generative modelling paradigm around the flexibility of the Kolmogorov equation.

The consideration of flows on discrete state spaces has also been used to construct GFlowNet algorithms Bengio et al. (2023) which aim to sample from a given energy function. Here we instead focus on the the generative modeling context where we aim to sample novel datapoints when only given access to some dataset of training examples. GFlowNets also can use the detailed balance equation Eq. 13 as a training training objective. Detailed balance is also used in Markov Chain Monte Carlo methods (Metropolis et al., 1953; Hastings, 1970) to construct a transition probability with the desired energy function that we wish to sample from as its stationary distribution. In our work, we use the detailed balance condition as a way to increase the inference time flexibility in our framework

Appendix E Multidimensional Data

In this section we derive how we can efficiently model $D$ dimensional data, $x_{1}\in\{1,\dots,S\}^{D}$ by using factorization assumptions. When we wish to emphasize the multidimensional aspect we can write $x_{1}^{1:D}$ and use $x_{1}^{d}\in\{1,\dots,S\}$ to refer to the value in dimension $d$. We use $1:D\backslash d$ to denote all dimensions except $d$. To operate in multidimensional spaces, we will make the following assumptions

Assumption 1 $p_{t|1}(x_{t}^{1:D}|x_{1}^{1:D})=\prod_{d=1}^{D}p_{t|1}(x_{t}^{d}|x_{1}^{d})$

Assumption 2 $p_{t|1}(x_{t}^{d}|x_{1}^{d})=0\implies\partial_{t}p_{t|1}(x_{t}^{d}|x_{1}^{d})=0,\forall d$

Assumption 3 $R_{t}(x_{t}^{1:D},j^{1:D}|x_{1}^{1:D})=\sum_{d=1}^{D}\delta\{x_{t}^{1:D\backslash d},j^{1:D\backslash d}\}R_{t}^{d}(x_{t}^{d},j^{d}|x_{1}^{d})$

The first assumption creates independent corruption processes in each dimension, similar to the factorization assumptions made in diffusion models where the forward noising processes proceed independently in each dimension. Assumption 2 is the same assumption we made in order to derive $R^{*}_{t}$ in $1$-dimension but now we assume it individually for every dimension. Finally, assumption 3 states that for our data conditional rate matrix, it decomposes into a sum of rate matrices for each dimension and so the rate for transitions that change more than $1$ dimension at a time are 0. This is the same assumption made by Campbell et al. (2022) in order to make calculations tractable. We will enable our process to make multiple dimensional changes simultaneously later when we come to derive our sampling algorithm.

Under these assumptions, we will now derive DFM for the multidimensional case. We start with the data conditional Kolmogorov equation

We now substitute the form for the rate matrix under Assumption 3 into the RHS of (158) to get

where we use $x_{t}^{1:D\backslash d}\odot j^{d}$ to denote a vector of dimension $D$ where in the $d$-th dimension it has the value of $j^{d}$ and in the other dimensions it has values $x_{t}^{1:D\backslash d}$. We now verify that the following form for $R_{t}^{d}$ satisfies the Kolmogorov equation,

where $\mathcal{Z}_{t}^{d}=|\{j^{d}:p_{t|1}(j^{d}|x_{1}^{d})>0\}|$ and we only define ${R^{*}_{t}}^{d}$ for off-diagonal entries, $x_{t}^{d}\neq j^{d}$ remembering that ${R^{*}_{t}}^{d}(x_{t}^{d},x_{t}^{d}|x_{1}^{d})=-\sum_{j^{d}\neq x_{t}^{d}}{R^{*}_{t}}^{d}(x_{t}^{d},j^{d}|x_{1}^{d})$.

We first assume $p_{t|1}(x_{t}^{d}|x_{1}^{d})>0\,\forall d$ and substitute in ${R^{*}_{t}}^{d}$ into equation (160).

where we have used the fact that $p_{t|1}(x_{t}^{1:D}|x_{1}^{1:D})=\prod_{d=1}^{D}p_{t|1}(x_{t}^{d}|x_{1}^{d})$.

For the case that there exists a $d^{\prime}$ for which $p_{t|1}(x_{t}^{d^{\prime}}|x_{1}^{d^{\prime}})=0$ we have $\partial_{t}p_{t|1}(x_{t}^{d^{\prime}}|x_{1}^{d^{\prime}})=0$ by assumption. We first examine the LHS of equation (158) in this case.

where we use $1:D\backslash d,d^{\prime}$ to mean all dimensions except $d$ and $d^{\prime}$. We now examine the RHS of equation (158).

where we have used the fact that ${R^{*}_{t}}^{d^{\prime}}(j^{d^{\prime}},x_{t}^{d^{\prime}}|x_{1}^{d^{\prime}})=0$ because $p_{t|1}(j^{d^{\prime}}|x_{1}^{d^{\prime}})=0$ Therefore, for both cases we have ${R^{*}_{t}}$ satisfies the conditional Kolmogorov equation (158) and thus we have found a rate matrix that generates our desired conditional flow. The final step is to convert this rate matrix conditioned on $x_{1}^{1:D}$ into an unconditional rate matrix that can be used for generative modeling. We first write down the unconditional multi-dimensional Kolmogorov equation

We now make the following assumption for the form of the unconditional rate matrix and verify that it indeed satisfies the unconditional multi-dimensional Kolmogorov equation, (181).

with ${R^{*}_{t}}^{d}(x_{t}^{d},j^{d}|x_{1}^{d})$ being given by (161). Substitute this form into (181)

where we have used Eq. 160 with the fact that we know ${R^{*}_{t}}$ given by Eq. 161 satisfies the conditional Kolmogorov equation Eq. 158. We have now verified that the rate given by Eq. 182 gives us our desired unconditional flow and we can use it for generative modeling.

In order to approximate the true generative rate matrix given by equation (182), we need approximations to the denoising distributions in each dimension, $p(x_{1}^{d}|x_{t}^{1:D})$, for $d=1,\dots,D$. We can parameterize these conditionally independent $x_{1}^{d}$ distributions through a neural network that outputs logits of shape $D\times S$ when given input $x_{t}^{1:D}$ of shape $D$. We then apply a softmax to the logits to obtain approximate denoising probabilities $p_{\theta}(x_{1}^{d}|x_{t}^{1:D})$, $d=1,\dots,D$ of shape $D\times S$. We learn the parameters of the neural network with the cross entropy loss for each dimension

E.2 Sampling

The standard Euler step transition probability for our CTMC defined through our learned denoising model with time step $\Delta t$ is

We note this sampling method is similar to the tau-leaping method used in prior CTMC based approaches Gillespie (2001); Campbell et al. (2022) however tau-leaping allows multiple jumps to be made in the same dimensions which is unsuitable for categorical data.

E.3 Detailed Balance

In this section we verify that if we achieve detailed balance individually and independently in each dimension, then our full dimensional process will also be in detailed balance.

Consider the multidimensional detailed balance equation

Now, substitute in our factorized forms for $R_{t}(x_{t}^{1:D},j^{1:D}|x_{1}^{1:D})$ and $p_{t|1}(x_{t}^{1:D}|x_{1}^{1:D})$

Now, both sides are for when $x_{t}$ and $j$ differ in more than one dimension. Consider the case when they differ in exactly one dimension, call it $d$. The detailed balance equation simplifies to

which we note is the standard single dimensional detailed balance equation for dimension $d$. Therefore, if our $R_{t}^{d}$ matrices are all in detailed balance with their respective $p_{t|1}(x_{t}^{d}|x_{1}^{d})$ conditional marginals, then the full dimensional rate matrix $R_{t}(x_{t}^{1:D},j^{1:D}|x_{1}^{1:D})$ will also be in detailed balance with the full dimensional conditional marginals $p_{t|1}(x_{t}^{1:D}|x_{1}^{1:D})$.

Appendix F Implementation Details

In this section we provide concrete derivations of our DFM method. We use a masking process in Sec. F.1, a uniform process in Sec. F.2 and explore the general case for any given $p_{t|1}$ in Sec. F.3. We also provide minimal PyTorch implementations for our training and sampling loops in each case. We will assume multi-dimensional data under the factorization assumptions listed in App. E.

Notebooks containing these minimal examples can be found at https://github.com/andrew-cr/discrete_flow_models.

Here, we assume the masking form for $p_{t|1}$. We begin by writing this data conditional flow

This is the distribution we will use to train our denoising model $p_{1|t}^{\theta}(x_{1}^{1:D}|x_{t}^{1:D})$. PyTorch code for the training loop is given in Listing LABEL:lst:masking_training

We will also derive the form for ${R^{*}_{t}}^{d}(i^{d},j^{d}|x_{1}^{d})$. For this we need to find $\partial_{t}p_{t|1}(x_{t}^{d}|x_{1}^{d})$.

We can now find ${R^{*}_{t}}^{d}(x_{t}^{d},j^{d}|x_{1}^{d})$. When working with rate matrices in this section, we will always assume $x_{t}^{d}\neq j^{d}$ and calculate the diagonal entries as $R_{t}(i,i)=-\sum_{j\neq i}R_{t}(i,j)$ later. We note that ${R^{*}_{t}}^{d}(x_{t}^{d},j^{d}|x_{1}^{d})=0$ for $p_{t|1}(x_{t}^{d}|x_{1}^{d})=0$ or $p_{t|1}(j^{d}|x_{1}^{d})=0$. Further, our initial distribution $p_{0}(x_{0}^{1:D})=\prod_{d=1}^{D}\delta\left\{x_{0}^{d},M\right\}$. Therefore, at all points in our CTMC, $x_{t}^{d}$ is only ever $M$ or $x_{1}^{d}$. Furthermore, we only ever have to consider transitions to a $j^{d}$ that is either $j^{d}=M$ or $j^{d}=x_{1}^{d}$. Now, for $p_{t|1}(x_{t}^{d}|x_{1}^{1:D})>0$ and $p_{t|1}(j^{d}|x_{1}^{1:D})>0$ we have

We note here that our calculation may not strictly be valid for exactly $t=0$ or $t=1$ but are valid for any $t\in(0,1)$ and so we can simply ignore these edge cases, see Sec. B.2 for further discussion. Now we find our unconditional rate matrix

Listing LABEL:lst:masking_sample shows PyTorch code that implements this sampling loop.

for some constants $a_{t}$ and $b_{t}$ that we must find. Substituting this into the detailed balance equation along with the masking interpolation form for $p_{t|1}(x_{t}|x_{1})$ gives

This equation must be true for all $i,j$. Pick $i=x_{1}$ and $j=M$ to get

We now combine this rate with ${R^{*}_{t}}^{d}$ that we calculated previously to find a new unconditional rate matrix with a variable amount of stochasticity.

where on the final line we have used the fact that $p_{1|t}^{\theta}(x_{1}^{d}=x_{t}^{d}|x_{t}^{1:D})=0$ for $x_{t}^{d}=M$ and $p_{1|t}^{\theta}(x_{1}^{d}=x_{t}^{d}|x_{t}^{1:D})=1$ when $x_{t}^{d}\neq M$ because if a dimension is unmasked then it must be the true $x_{1}$ value under our definition of $p_{t|1}(x_{t}|x_{1})$. We now find our transition probabilities

where again we have used the fact that $p_{1|t}^{\theta}(x_{1}^{d}=M|x_{t}^{1:D})=0$. Inspecting $p_{t+\Delta t|t}(j^{d}|x_{t}^{1:D})$ for $j^{d}\neq x_{t}^{d}$, we see that if $x_{t}^{d}=M$ then we have an overall probability of unmasking of $\frac{1+\eta t}{1-t}\Delta t$ and once we do unmask, the new value is drawn from $p_{1|t}^{\theta}(x_{1}^{d}|x_{t}^{1:D})$. This is like before but now there is a bonus probability of unmasking of $\frac{\eta t}{1-t}$. When $x_{t}^{d}\neq M$ then we have a probability of $\eta\Delta t$ of jumping back to the mask state. This creates a flux of states switching back and forth between masked and unmasked for $\eta>0$ hence why these processes are more ‘stochastic’. However, because when $\eta$ is increased we also increase the rate at which we unmask, the desired conditional flow $p_{t|1}(x_{t}|x_{1})$ is maintained for any value of $\eta$. Listing LABEL:lst:masking_sample_with_noise shows PyTorch code that implements sampling with this extra stochasticity.

Our method has similarities to other discrete diffusion models when using this form for $p_{t|1}$ and we clarify these links in Sec. H.2.

F.1.2 Purity Sampling

When using the masking form for $p_{t|1}$ we can also easily implement a purity sampling scheme Tang et al. (2022). This sampling method decides which dimensions to unmask based on an estimate of the model confidence in that dimension’s final value. Currently, our sampling method will uniformly at random choose which dimension to unmask. To improve upon this approach, purity sampling will instead rank dimensions based on which dimension has the highest model probability. More specifically, for each dimension we calculate a purity score for dimension $d$ defined as

For the next simulation step, we then decide how many dimensions should be unmasked. The number of dimensions to unmask is binomially distributed with probability of success $\frac{\Delta t}{1-t}$ and number of trials equal to the number of dimensions that are currently masked. Once we have sampled a number of dimensions to unmask from this binomial distribution, we then unmask that number of dimensions starting from the dimension with highest purity score, then the dimension with second highest purity score and so on. We only consider dimensions that are currently masked to be eligible for unmasking. When using $\eta>0$, the probability of success in our binomial distribution increases to $\Delta t\frac{1+\eta t}{1-t}$ and so on average more dimensions get unmasked during each simulation step. At the end of each simulation step, we then remask a sample of randomly chosen dimensions which are uniformly chosen at random each with a probability $\Delta t\eta$ of being chosen.

F.2 Uniform Example

In this section we walk through the derivation and implementation of DFM when using the uniform based interpolation distribution. We start with the data conditional marginal distribution

This distribution is all that is needed to train the denoising model $p_{1|t}^{\theta}(x_{1}^{1:D}|x_{t}^{1:D})$. We give PyTorch code for the training loop with the uniform interpolant in Listing LABEL:lst:uniform_training.

In order to sample our trained model, we will need to derive ${R^{*}_{t}}^{d}(i^{d},j^{d}|x_{1}^{d})$. The first step is to find $\partial_{t}p_{t|1}(x_{t}^{d}|x_{1}^{d})$,

We will now find ${R^{*}_{t}}^{d}(x_{t}^{d},j^{d}|x_{1}^{d})$. As before we will always assume $x_{t}^{d}\neq j^{d}$ and calculate diagonal entries as needed using the relation $R_{t}(i,i)=-\sum_{j\neq i}R_{t}(i,j)$.

The only non-zero value is when $j^{d}=x_{1}^{d}$ and $x_{t}^{d}\neq x_{1}^{d}$ and so ${R^{*}_{t}}^{d}(x_{t}^{d},j^{d}|x_{1}^{d})$ is

We can now find the unconditional rate matrix, still assuming $x_{t}^{d}\neq j^{d}$

Listing LABEL:lst:uniform_sample shows PyTorch code that implements this sampling loop.

Now, this equation must be true for any $i\neq j$. Pick $i=x_{1}$ and $j\neq x_{1}$ to get

This can now be combined with ${R^{*}_{t}}^{d}$ to create a new unconditional rate matrix with a variable amount of stochasticity.

We can interpret this rate, with the first term being the rate at which we should transition to states that are predicted to correspond to the clean data. The second term is a ‘noise term’ which creates transitions away from the current state if it is predicted to correspond to the final clean data. The first term then has additional weighting as $\eta$ is increased to counter act this effect. The effect of the stochasticity is then to create a flux going on and off the predicted final clean state during generation. We now find our transition probabilities

We can find $p_{t+\Delta t|t}(j^{d}|x_{t}^{1:D})$ for $j^{d}=x_{t}^{d}$ programmatically as before by requiring that the probability vector sum to $1$. Listing LABEL:lst:uniform_sample_with_noise shows the implementation for the uniform interpolant with noise.

F.3 General Case

We now describe the training and sampling loop for a general conditional flow $p_{t|1}(x_{t}|x_{1})$. We require this interpolant to be factorized, $p_{t|1}(x_{t}^{1:D}|x_{1}^{1:D})=\prod_{d=1}^{D}p_{t|1}(x_{t}^{d}|x_{1}^{d})$, be differentiable and have $p_{t|1}(j^{d}|x_{1}^{d})=0\implies\partial_{t}p_{t|1}(j^{d}|x_{1}^{d})=0$. We assume that we have access to functions that can sample from $p_{t|1}(x_{t}|x_{1})$, evaluate $p_{t|1}(x_{t}|x_{1})$ and evaluate $\partial_{t}p_{t|1}(x_{t}|x_{1})$. Our training loop consists of sampling data, sampling $x_{t}\sim p_{t|1}(x_{t}|x_{1})$ and training with the cross entropy loss, see Listing LABEL:lst:general_training.

Appendix G CTMC Sampling Methods

The form of this update means that we don’t necessarily need to calculate the full expectation over $R_{t}(x_{t},j|x_{1})$. We can simply sample $x_{1}$ from $p_{1|t}^{\theta}(x_{1}|x_{t})$ and then plug this sample into $R_{t}(x_{t},j|x_{1})$ which we then use in the Euler update. To see that this strategy still samples from the same distribution over $x_{t+\Delta t}$, we can write the distribution over $x_{t+\Delta t}$ as $p_{t+\Delta t|t}$,

and so $p_{t+\Delta t|t}(x_{t+\Delta t}|x_{t})$ can be seen as the marginal of joint distribution $p_{1|t}^{\theta}(x_{1}|x_{t})p_{t+\Delta t|t}(x_{t+\Delta t}|x_{1},x_{t})$. Therefore, to produce a sample $x_{t+\Delta t}$ from $p_{t+\Delta t|t}(x_{t+\Delta t}|x_{t})$, we can instead sample $x_{1},x_{t+\Delta t}$ from the joint distribution $p_{1|t}^{\theta}(x_{1}|x_{t}^{1:D})p_{t+\Delta t|t}(x_{t+\Delta t}|x_{1},x_{t})$, and take only the $x_{t+\Delta t}$ part of this joint sample.

Another method to simulate a CTMC is $\tau$-leaping, Gillespie (2001); Campbell et al. (2022) which allows multiple jumps to be made both across dimensions and within each dimension. Multiple jumps within a single dimension does not make sense for categorical data where there is no ordering, however, it can be useful for ordinal data such as a discretized image where the $\tau$-leaping update allows multiple jumps to be applied at once to cover a larger distance. To calculate a $\tau$-leaping update, a Poisson random variable needs to be drawn with the rate matrix giving the rate parameter. Therefore, for this type of update, the full unconditional $R_{t}^{\theta}(x_{t},j)$ would need to be calculated.

We finally note that there is a body of work creating CTMC samplers for generative models (Sun et al., 2023b; Lou et al., 2023) that may be faster to simulate than the standard Euler step. In this work, we focus on framework simplicity, not optimizing for sampling speed and leave application of these approaches as future work.

Appendix H Comparison with Discrete Diffusion Models

In this section we clarify the relationship between DFM and classical discrete diffusion models. In Sec. H.1 we compare to continuous time models using the uniform corruption process as an example. In Sec. H.2 we compare to discrete time models using the masking process as the example.

Here we compare to continuous time discrete diffusion models (Campbell et al., 2022) using the uniform corruption process as an example. In this section, we will assume $t=0$ is pure noise and $t=1$ is clean data which we note is a flipped definition of time to classical diffusion models to aid in our comparison with DFMs.

To calculate $p_{t|1}(x_{t}|x_{1})$ we can use the equation

where $(P_{t})_{ij}=p_{t|1}(x_{t}=j|x_{1}=i)$ and $\exp$ is the element wise exponential. By the symmetry of the problem, we can infer that $p_{t|1}(x_{t}=j|x_{1}=i)$ will have only two possible values. Either $j=i$ and we are finding the probability of staying at $i$, or $j\neq i$ and we are finding the probability of having left $i$, and since uniform corruption treats all states equally, these will be same quantities for any starting state and any state $j\neq i$. So to find our schedule we just need to consider one element of the matrix $P_{t}$. Let us consider an off-diagonal element $i\neq j$ of $P_{t}$, which will have probability

We will try and match this to the simple linear schedule that we have had as our running example in the explanation of DFM.

The next step in a discrete diffusion model is to find the time reversed rate matrix $\hat{R}_{t}$ which gives a CTMC that runs in the opposite direction to $R_{t}$ and generates novel data from noise. Here $\hat{R}_{t}$ is running from time $t=0$ at noise towards clean data at $t=1$. From Campbell et al. (2022), we have

We notice a similarity to the DFM equations, where the generative rate is the expectation of a quantity with respect to $p_{1|t}(x_{1}|i)$. Indeed we now show that $R_{t}(j,i)\frac{p_{t|1}(j|x_{1})}{p_{t|1}(i|x_{1})}$ is a $x_{1}$ conditioned rate matrix $R^{\text{diff}}_{t}(i,j|x_{1})$ that achieves the conditional flow $p_{t|1}(i|x_{1})$. Consider the Kolmogorov equation

Substitute in our form for $R^{\text{diff}}_{t}$

where on the second to last line we have used the fact that the corruption matrix $R_{t}(i,j)$ when started at $p_{t=1}(x_{t}|x_{1})=\delta\left\{x_{t},x_{1}\right\}$ will evolve the marginals according to $p_{t|1}(x_{t}|x_{1})$ because this is how we derived $p_{t|1}(x_{t}|x_{1})$ in the first place. Note $R_{t}$ runs in the reverse direction hence the negative sign.

with $b_{t}$ and $c_{t}$ being fully flexible (provided they are positive). Using the form for $R^{*}_{t}(i,j|x_{1})=\frac{1}{1-t}\delta\left\{j,x_{1}\right\}(1-\delta\left\{i,x_{1}\right\})$ that we derived in Appendix F.2 we have

which is equal to $R^{\text{diff}}_{t}$ if we have $b_{t}=c_{t}=\frac{1}{3t}$.

In summary, we have found that classical discrete diffusion models make an implicit choice for $R_{t}(i,j|x_{1})$ which corresponds to a certain level of stochasticity in the CTMC and that the choice is made at training time because the rate matrix is used in the ELBO objective. Further, we have seen it is much harder to derive the noise schedule $p_{t|1}(x_{t}|x_{1})$ in classical discrete diffusion models due to the need to be able to apply the matrix exponential to $R_{t}$. In DFM, we can simply write down the $p_{t|1}(x_{t}|x_{1})$ noise schedule we want and we are not restricted in having to pick $R_{t}$ that are amenable to matrix exponentiation. We also get to choose any $R_{t}(i,j|x_{1})$ at test time rather than being fixed to the implicit choice of $R^{\text{diff}}_{t}$.

H.2 Discrete Time Discrete Diffusion Models

In this section we will clarify the link to the discrete time diffusion method D3PM (Austin et al., 2021) when using the masking process for both methods. Here, we will use the convention from Austin et al. (2021) of using $t=0$ for clean data and $t=T$ for noise.

We will first summarize the key results from Austin et al. (2021) when using the absorbing state process which is a different name for a masking type process (the mask is the absorbing state). $t$ can take on any discrete value in $t\in\{0,1,\dots,T\}$. The diffusion model is first defined using a noising transition kernel

From this transition kernel, we can then calculate the noise marginals, $p(x_{t}|x_{0})$

We then define our generative reverse process as

Continuing with the D3PM example using the absorbing state process, we obtain the following form for $p_{\theta}(x_{t-1}|x_{t})$

When we set $\beta_{t}=\frac{1}{T-t+1}$, we obtain a linear noise schedule giving

Now, let us define $\xi\vcentcolon=\frac{t}{T}$ to be the proportion that the process is through the total number of time steps. $\xi\in$ and if we consider it to be an analogue of our continuous time variable, we can see that the original discretization steps of D3PM correspond to a discretization of the $$interval with timesteps of$\Delta t=\frac{1}{T}$. Substituting these definitions into our update step gives,

Appendix I Text Experiment Details

Code for our text experiments can be found at https://github.com/andrew-cr/discrete_flow_models.

For our denoising network we use the transformer architecture Vaswani et al. (2017) as implemented in the nanoGPT repository, https://github.com/karpathy/nanoGPT. We generally follow the smallest GPT2 architecture Radford et al. (2019). At the input we have our input tokens $x_{t}$ of shape $B,D$ where $B$ is the batch size and $D$ is the number of dimensions i.e. the sequence length, our time $t$ of shape $B$, and, if we are self-conditioning, prior $x_{1}$ prediction tokens of shape $B,D$. We embed the $x_{t}$ and $x_{1}$ tokens using the same learned embedding, and use a model embedding size of $768$ resulting in tensors of shape $B,D,768$. We embed the position of each token using a learned embedding for each possible position. We embed the time $t$, using Transformer sinusoidal embeddings following Ho et al. (2020). We train all our diffusion models with self-conditioning Chen et al. (2023). To input the prior $x_{1}$ prediction, we stack the $x_{t}$ embedded tensor $B,D,768$ with the $x_{1}$ prior prediction token tensor $B,D,768$ to obtain a tensor of shape $B,D,768\times 2$. We then apply a linear layer to project down to the model embedding dimension resulting in a tensor of shape $B,D,768$. Before applying transformer blocks, we add together the $x_{t}$ (and $x_{1}$) embedding tensor, the position embedding and the time embedding to obtain the final $B,D,768$ input tensor.

The transformer stack consists of $12$ transformer blocks, each block consisting of a LayerNorm, SelfAttention, LayerNorm, MLP stack. Within our SelfAttention block, we use $12$ heads and apply Qk-layernorm Dehghani et al. (2023) to our query and key values as we observed this improved convergence. Our MLP blocks consist of a $768\rightarrow 768\times 4$ linear layer, followed by a GELU activation, followed by a $768\times 4\rightarrow 768$ linear layer. We do not apply dropout. Our output layer consists of a linear head with output dimension $28$. We use $28$ token categories, $26$ lower case letters, a whitespace character and a mask token. The model outputs logits of shape $B,D,28$ which we then apply a softmax to, to obtain $p_{\theta}(x_{1}|x_{t})$ probabilities.

The dataset text8 is $100$MB of text data from English Wikipedia. The text is all converted to lower case letters, i.e. capital letters are converted to lower case and numbers are written as text, i.e. $8$ becomes ‘eight’.

During training, we use a batch size of $256$ with $8$ gradient accumulation steps. We train on sequences of length $256$. The model is therefore trained on $524,288$ tokens per gradient update. To train self-conditioning, on $50\%$ of training iterations, we input prior $x_{1}$ prediction tokens as all masks so that the model learns to be able to predict $x_{1}$ without any prior information. On the other $50\%$ of training iterations, we perform two model forward passes. We first predict $x_{1}$ using masks as the prior $x_{1}$ tokens to obtain an initial set of $p_{\theta}(x_{1}|x_{t})$ logits. We then sample from the initial $p_{\theta}(x_{1}|x_{t})$ distribution to obtain predicted $x_{1}$ tokens. We then feed these tokens back into the model through the self-conditioning input and predict the $x_{1}$ logits once more. These logits are then used in the loss. We only back propagate through the second forward pass of the model.

When training the D3PM model, we found that the default cross entropy weighting of $1/t$ (with a flipped definition of time) resulted in poor convergence and so we applied an equal weighting of the cross entropy across time to be consistent with the DFM loss.

We train our D3PM and DFM models for $750$k iterations on 4 Nvidia A40 GPUs using a learning rate of $10^{-4}$ and $1000$ linear warm up steps. We use a cosine decay schedule after the initial warm up towards a minimum learning rate of $10^{-5}$ which would be reached at $1$M iterations. We use the AdamW optimizer Loshchilov & Hutter (2017) with weight decay parameter $0.1$. We monitor the validation loss throughout training. Validation loss continues to drop throughout training and we evaluate the final $750k$ model in our experiments. When training the autoregressive model, we use the same architecture but find that it begins to overfit the data much faster than the diffusion based models. After $3500$ iterations the validation loss begins to increase and so we use the model with minimum validation loss in our evaluations. This is consistent with findings that autoregressive models require much less compute to converge than diffusion based models Gulrajani & Hashimoto (2023).

We use the masking interpolant in our DFM with linear interpolant, as described in Appendix F.1. For D3PM, we use the absorbing state corruption process, the links to the DFM process are described in Appendix H.2.

For evaluation, we sample the DFM with $\Delta t=0.001$. We simulate up to $t=0.98$ and then for any remaining tokens that are still mask, we set them to the most likely token under the model’s denoising distribution, $p_{\theta}(x_{1}|x_{t})$. We stop simulating at $t=0.98$ to avoid any singularities similar to how diffusion models stop near $t=0$. For D3PM we train with $1000$ timesteps to match DFM.

For each temperature setting applied to the $p_{\theta}(x_{1}|x_{t})$ logits, we sample $512$ sequences all of length $256$ tokens. We then calculate the negative log-likelihood assigned to each sequence using GPT-J-6B Wang & Komatsuzaki (2021) and the BPE tokenizer Radford et al. (2019). We then average the negative log-likelihoods over the $512$ sequences. The sample entropy is calculated by first tokenizing with the BPE tokenizer and then calculating the entropy as $\sum_{i}-p_{i}\log p_{i}$ where $p_{i}$ is the empirical probability of token $i$ estimated using the full set of $512$ samples. Tokens for which $p_{i}=0$ are not included in the sum. For reference, the dataset achieves a negative log-likelihood of 4.2 as measured by GPT-J-6B.

Here we examine the effect of the noise level $\eta$ on the sample quality of generations from our DFM method. We follow the follow the same procedure as before but vary $\eta$ with values $\eta=0,1,2,5,10,15,20,30,50$. We plot the results in Figure 4. We find that generally, as the noise level increases, we lower our negative log-likelihood. However, we find that if the noise level is increased too much, then degenerate behaviour can occur, for example when $\eta=50$, at high logit temperatures the negative log-likelihood increases and the sample entropy decreases away from the dataset. Observing the samples, we find that the model generates incoherent text at this point. We find that the intermediate noise level $\eta=15$ provides good sample quality whilst avoiding this behaviour.

I.2 Example Text Generations

In this section we provide non cherry picked generations from the text models. For each model we have swept over the temperature applied to the logits and it would be impractical to include examples for all models for all temperature settings. Instead, we select one temperature setting for each model such that the samples have similar entropy but vary in negative log-likelihood. We show the selected temperature settings in Figure 5.

D3PM Temperature 0.8 Samples: ved as a personal area to form the five counties of the area and a country with their own which is usually called paris gietgothic can also lead an area to work in divisions over a pileur as in the name of man the bears have over the last two years from th one five zero zero zero zero press money to present this to a meschasel linear industrial base ulse sudan expanded its economy and accounts for car prices and two eight five more than one zero zero of the largest industrial inventions over the world were eed alternatively as human being and the anti constitutionalay doctrines a particular example of the concept is one reason for human rights or as in certain regions there is a double constitution more recognized region of europe in this region the glass an DFM $\eta=0$ Temperature 0.8 Samples: ed era vol seven one nine one one december one nine six one junju that s one of nine one one country page of love footnote pages charles s feadman history of the red sea corea one nine nine one red sea vol one january one nine nine seven flying profiles ch allowes the vectores to be composed as systems of data for example no machine is a computer one would do not know where there are undirected storage of other data storage particularly the computer science eve to substitute such a based data that is one of me io the plate n and feminine along the trail to change the amount of naturated information in the start tape selective figurative memory the mind is determined by the second net on the string c with two buttons the tag retes the header when queued the se Autoregressive Temperature 0.9 Samples: licklyn american football coach to holy roman emperor and roman stories radio and facilities in the u s civil rights movement the dc circuit collection of the witches leading the transissario times and spinoffs to american cartoonists cartoonist kyle marci the british one one eight four minamoto minister or al di nortello ministries son of monte oise klepe which chose to give up its character on the go he was known to publish a wade of white performances started in one eight five one kleine married the gigan mausoleum in one eight one six alabama was engaged by a large scale as we know alabama migration the palace of westminsters and proceeded to father she also learned to speak with the abramic mouth of the space the replica was apparently built de provence g DFM $\eta=15$ Temperature 0.8 Samples: e curous greek by alexander van hep ven see archaic origin of the word cupola another meaning suggests that the word kupola is the latin word cupei kupolum old german derived from the latin word for the river the name comes from a latin word for tree with es so balloonists refine this combination specifically to preserve your own land in the runner both examples of clean steering creating agout like rods that produced successful rods and for the end the first few pistols compact stunt a musical setting mult by reign over agassi is considered a greatest match by the day he will never play and will continue to be imitated agassi can play determinedly but agassi would always look to the victorious build he should not finish years going up to then that he would b

Appendix J Protein Generation Experiment Details

We present additional experiment details and results for protein generation with Multiflow.

Code for Multiflow and experiments can be found at https://github.com/jasonkyuyim/multiflow

We use an architecture modified from the FrameDiff architecture from Yim et al. (2023b). This architecture consists of Invariant Point Attention (Jumper et al., 2021) combined with transformer blocks, we refer to Yim et al. (2023b) for in-depth details. We modify this network architecture by increasing the number of network blocks to 8, increasing the number of transformer layers within each block to 4, decreasing the number of hidden channels used in the IPA calculation to 16, removing skip connections and removing psi-angle prediction. To enable our model to output logits for the discrete $p_{1|t}^{\theta}(x_{1}|x_{t})$ distribution, we add an output 3 layer MLP with the same embedding size as the main trunk. This results in a network with 21.8M parameters.

In Yim et al. (2023b), psi-angle prediction is used to infer the location of oxygen atoms, however, this position can be inferred to high accuracy using prior knowledge of the backbone structure of proteins, following Yim et al. (2023a).

J.2 Additional Multiflow Results

We show results of Multiflow across more lengths than done in Sec. 6.2.1 and show that using the ESMFold oracle for data distillation still gives improved performance when we switch the evaluation oracle to AlphaFold2.

Our results in Sec. 6.2.1 only evaluated 4 lengths (70, 100, 200, 300) to match the benchmark in RFdiffusion. However, other works have evaluated designability across all the lengths the method was trained on. We follow Protpardelle (Chu et al., 2023) to use Multiflow in generating 8 samples per length in the range $\{$50, 51, $\dots$, 400$\}$. Fig. 6 shows the results in the same format as Figure 2B in Protpardelle. We see Multiflow achieves near perfect designability up to around length 350 at which point designability starts to drop. This is expected since Multiflow was only trained on lengths up to 384, but also demonstrates the ability to generalize beyond the lengths it was trained on. We see Multiflow also achieves a desirable spread of secondary structure. We show samples above length 370 with the highest and lowest Co-design 1 RMSD in Fig. 7.

AlphaFold2 evaluation oracle. In Sec. 6.2.1, we presented a distillation technique of filtering out training examples that did not pass the designability criterion. This also involved adding more proteins to the training set after sampling structures with Multiflow and filtering with designability using ProteinMPNN and ESMFold. A potential risk of distillation is our model may overfit to ESMFold since this model is used to filter training data and also for evaluation. We show this is not the case in Table. 5 by presenting the Co-design 1 results using AlphaFold2 (AF2) as an alternative oracle. Our main results do not use AF2 since it is very slow and cumbersome to run and evaluate all our baselines. We evaluated Multiflow with and without distillation to test if distillation with ESMFold provides an improvement regardless of the oracle used at evaluation. Overall designability numbers are lower with AF2; however, in both columns we see there is a two fold improvement regardless of the evaluation oracle. This demonstrates distillation is not overfitting to the oracle used at evalution.

J.3 Uniform Conditional Flow Ablation

We ablate our use of the masking conditional flow and train a version of our Multiflow model using the uniform conditional flow ( see Sec. F.2). We assessed the model’s co-design performance by measuring the Co-Design 1 designability and diversity versus stochasticity level used at inference time. We also measure the secondary structure composition of the generated samples versus stochasticity level. Our results are given in Fig. 8. We find that in general, the Co-Design 1 designability increases with increasing stochasticity whilst the diversity as measured by the number of structural clusters decreases. We can see the reason when examining the secondary structure statistics versus stochasticity. We see that at high stochasticity levels, the model heavily favours generating alpha helices at the expense of beta strands thus reducing the overall structural diversity. This will be due to interactions between errors in the model and the ‘churn’ induced by extra stochasticity. It may be counter-intuitive that extra stochasticity reduces model diversity however we hypothesize that this is linked to the stochasticity inducing the model to converge on local optima in the likelihood landscape. When the model is generating a sample that it is confidence about, extra stochasticity will not shift it away from continuing down this simulation trajectory. However, when the model is exploring lower likelihood regions, the stochasticity can shift the models trajectory until it becomes stuck in a local optima again.

We find an overall worse trade-off between diversity and designability when using the uniform interpolant and so opt to use the masking interpolant in our main models.

J.4 Forward and Inverse Folding Experiments

The goal of our work is to develop the missing piece for a general-purpose framework for protein generation – namely DFM to integrate discrete data generation with a flow model. We combined DFM and FrameFlow to develop Multiflow where we have flexibility at inference time to choose which modality to provide and which to generate. The task we focus on in this work is co-generation where the structure and sequence are jointly sampled rather than one after the other as done in prior works. The other useful tasks in protein modeling are forward and inverse folding. The two tasks are briefly described as follows; more in-depth description can be found in Gao et al. (2020).

Forward folding: the task is to take the sequence as input and predicts the most thermodynamically plausible structure of the sequence. During evaluation, the ground truth structure is known, so we calculate the aligned structure erorr between the prediction and the ground truth. Several metrics exist to compute structure error, such as the Global Distance Test (GDT) commonly used in biophysical modeling (Pereira et al., 2021). We choose to use the aligned backbone RMSD error to keep our analysis simple and intuitive. The most well-used methods are AlphaFold2 (Jumper et al., 2021), RosettaFold (Baek et al., 2021), and ESMFold (Lin et al., 2023). AlphaFold2 and RosettaFold rely on using evolutionary information which our model does not have access to (though can be extended to use). We compare against ESMFold, which does not use explicit evolutionary information, and due to its speed.

Inverse folding: the task is to use the structure as input and predict the most likely sequence that would forward fold into the structure. By this definition, the most sensible metric is the designability metric also used for co-generation. Specifically, the inverse folding model generates a sequence and we use ESMFold to predict the structure given this generated sequence. We call the self-consistency RMSD (scRMSD) as the RMSD between the structure predicted by ESMFold and the original input structure (Trippe et al., 2022). The objective is to minimize scRMSD. The de facto method for inverse folding is ProteinMPNN (Dauparas et al., 2022). Hence we compare against ProteinMPNN.

It is important to emphasize that different deep learning models have been specifically developed for forward and inverse folding, but no method can accomplish both tasks nor co-generate both sequence and structure. Multiflow is unique in this regard to be able to perform co-generation, forward folding, and inverse folding. We leave improving forward and inverse folding performance as a future work. Our aim is to demonstrate baseline performance of using a co-generation method to perform forward and inverse folding. We hope others can aid in advancing general purpose protein generative models.

ESMFold and ProteinMPNN have their own training and test sets which makes rigorous comparison impossible. Re-training ESMFold and ProteinMPNN with the same training set of Multiflow is beyond the scope of our work. Our results are a initial baseline of how Multiflow generally fares to specialized models on forward and inverse folding.

Our test set is based on a time-based split of the PDB. We downloaded structures and sequences from the PDB that were released between 1st September 2021 and 28th December 2023. This time-based split ensures that none of the test set proteins are present in the training data for Multiflow, ProteinMPNN or ESMFold. We then select all single chain monomeric proteins with length between 50 and 400 inclusive. We further filter out proteins that are more than 50% coil residues and proteins that have a radius of gyration in the 96th percentile of the original dataset or above. We also filter out structures that have missing residues. We cluster proteins using the 30% sequence identity MMSeqs2 clustering provided by RCSB.org. We take a single protein from each cluster that matches our filtering criteria. This gives us a test set of 449 proteins with minimum length 51 and maximum length 398.

J.4.1 Forward Folding Results

In Fig. 9 we examine the distribution of errors on our test set for both ESMFold and Multiflow. We find that generally Multiflow can have some success with proteins of smaller length but struggles with longer proteins. We investigate salient test examples from the plot to understand success and failure modes of our model. Multiflow is generally able to predict realistic protein structures with often similar secondary structure distributions as to the ground truth example seen by having similar proportions of non-loop residues between the ground truth and predicted structure. However, Multiflow often fails to predict the exact folded structure with high accuracy.

We quantify the secondary structure prediction accuracy in Fig. 10 by comparing the secondary structure present in the ground truth versus the structure predicted by Multiflow. We find good correlation between the predicted secondary structure and ground truth highlighting that Multiflow is able to use information present within the given sequence to generate structures.

J.4.2 Inverse Folding Results

We plot our results in Fig. 11. We find that Multiflow performs competitively with PMPNN across a wide range of protein lengths with PMPNN achieving slightly lower scRMSD values on average. For both models, scRMSD tends to cluster around 1 to 2 scRMSD. There are test proteins for which PMPNN achieves a lower scRMSD and also cases protein for which Multiflow acheives the lower scRMSD.