Source-linked AI summary

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

Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, Tommi Jaakkola

arXiv:2402.04997v2stat.MLcs.LGq-bio.QM

TL;DR

Generative models need to handle discrete and continuous data, but flow-based models lacked a discrete counterpart for multimodal tasks. This paper introduces CTMC-based Discrete Flow Models and reports state-of-the-art protein co-design performance with a model that can condition on either sequence or structure.

  • Problem

    Flow-based multimodal generation lacked a discrete-data model, despite the need to model continuous protein structures and discrete amino acid sequences jointly.

    Source: §1 · p. 1
  • Method

    Discrete Flow Models use Continuous Time Markov Chains for discrete generation and combine with continuous-state flows to jointly generate protein structures and sequences.

    Source: §1 · pp. 1–2
  • Results

    Multiflow achieves state-of-the-art protein co-design performance and can jointly generate sequence and structure while conditioning on either modality.

    Source: §1 · p. 2
  • Takeaways & Limitations

    CTMC stochasticity enables control over generated samples’ secondary-structure composition and diversity.

    Source: §1 · p. 2

    The paper notes that general-purpose protein modeling could be misused to design toxic substances and omits specific methods readily adapted to harmful design.

    Source: Impact statement · p. 10

Abstract

from arXiv · show

Combining discrete and continuous data is an important capability for generative models. We present Discrete Flow Models (DFMs), a new flow-based model of discrete data that provides the missing link in enabling flow-based generative models to be applied to multimodal continuous and discrete data problems. Our key insight is that the discrete equivalent of continuous space flow matching can be realized using Continuous Time Markov Chains. DFMs benefit from a simple derivation that includes discrete diffusion models as a specific instance while allowing improved performance over existing diffusion-based approaches. We utilize our DFMs method to build a multimodal flow-based modeling framework. We apply this capability to the task of protein co-design, wherein we learn a model for jointly generating protein structure and sequence. Our approach achieves state-of-the-art co-design performance while allowing the same multimodal model to be used for flexible generation of the sequence or structure.

1. Introduction

The introduction presents Discrete Flow Models (DFMs), which use Continuous Time Markov Chains to extend flow-based generation to discrete data, and applies them in Multiflow for joint protein sequence and structure generation. It motivates this approach by the need for flexible multimodal generation and reports promising protein co-design results.

1. Introduction

  • Protein co-design is motivated by the interplay between structure, which endows protein function, and sequence, which provides the blueprint for how structure is made.The task jointly generates continuous protein structures and corresponding discrete amino-acid sequences.Source: §1 · p. 1
  • The introduction identifies a gap: diffusion spans discrete and continuous spaces but has inflexible sampling, while flow-based models offer sampling flexibility but lack a discrete-space formulation.The paper first tests DFMs on small-scale text data and reports that their expanded sample-time flexibility outperforms D3PM (Austin et al., 2021).Source: §1 · pp. 1–2
  • DFMs realize discrete flow through CTMC trajectories that follow a noise-to-data probability flow, with a denoising neural network trained using cross-entropy.CTMC stochasticity controls the number of trajectory transitions and can be adjusted at inference, whereas discrete diffusion fixes it at training time.Source: §1 · pp. 1–2
  • Combining a DFM for sequences with continuous flow-based structure generation, Multiflow jointly generates protein sequence and structure and can condition on either modality.Earlier multimodal approaches generated one modality and used a prediction model to infer the other.Source: §1 · p. 2
  • Multiflow achieves state-of-the-art protein co-design performance, while data distillation enables state-of-the-art structure generation; preliminary folding results suggest potential as a general-purpose protein model.The introduction also reports that CTMC stochasticity can control sample properties such as secondary-structure composition and diversity.Source: §1 · p. 2

2. Background

The background defines discrete sequences over a finite state space and introduces CTMCs as a framework for describing their trajectories and probability flows.

2. Background

  • A sequence has D dimensions, each taking one of S states; the exposition assumes D = 1, while results extend to D > 1 as discussed in App. E.Source: §2 · p. 2

2.1. Continuous Time Markov Chains.

  • A CTMC trajectory alternates between resting in a state and jumping, with the rate matrix determining jump frequency and destination; finite intervals can be simulated with Euler steps.The rate matrix and initial distribution together define the CTMC.Source: §2.1 · pp. 2–3

2.2. Kolmogorov equation

  • The Kolmogorov equation states that the marginal’s time derivative equals R_t^⊤p_t, defining an ODE whose distributions over time form a probability flow.A CTMC is defined by an initial distribution p0 and rate matrix R_t; its probability flow comprises the marginals of its trajectory at each time.Source: §2.2 · p. 3

3. Discrete Flow Models

Discrete Flow Models (DFMs) define a probability flow from noise to data and sample trajectories from it using continuous-time Markov chains (CTMCs). Their construction supports combining discrete and continuous flow models for multimodal generation.

3. Discrete Flow Models

  • DFMs define a discrete-data probability flow from noise to data, then simulate trajectories that follow it; this construction can combine with continuous-data flows for multimodal generation.Source: §3 · p. 3

3.1. A Flow Model for Sampling Discrete Data

  • The generative flow is built from explicit conditional flows that interpolate from noise to a datapoint, beginning at the noise distribution and converging to that datapoint at t = 1.Source: §3.1 · p. 3

3.1.1. SAMPLING

  • The sampling rate matrix is an expectation of datapoint-conditional rate matrices under the posterior over clean data, approximated by a neural denoiser to simulate trajectories from noise to data.Source: §3.1.1 · p. 3

3.1.2. TRAINING

  • The denoiser is trained with cross-entropy to predict clean datapoints from noisy samples, and the loss is independent of the conditional rate matrix, allowing that choice to remain flexible at inference.Source: §3.1.2 · p. 4

3.2. Choice of Rate Matrix

  • Given the zero-mass-state condition, a valid base rate can be augmented with detailed-balance rates without changing the flow; η tunes CTMC stochasticity, while η = 0 minimizes expected jumps under the stated factorization assumptions.Source: §3.2 · pp. 4–5

3.3. DFMs Recipe

  • The DFM recipe defines a conditional noise schedule, trains a denoiser, chooses a rate matrix, and runs sampling.Source: §3.3 · p. 5

4. Multimodal Protein Generative Model

Multiflow combines a discrete flow model for amino-acid sequences with a continuous flow model for protein structures, jointly generating both modalities in one model.

4.1. Multimodal Flows

  • Each residue is represented by backbone coordinates in SE(3) and an amino acid or mask state; side-chain atoms are left for future work.The residue state combines a Carbon-α translation, a local-frame rotation, and one of 20 amino acids or the mask state.Source: §4.1 · p. 5
  • The conditional flow factorizes across modalities and residues, using continuous ODEs for translations and rotations and a CTMC for amino acids; Proposition 4.1 verifies the resulting flow.The unconditional generative flow couples modalities and residues because its velocities and rate matrix depend on the entire corrupted protein state.Source: §4.1 · pp. 5–6

4.2. Training

  • Training independently samples structure and sequence noise levels, enabling learning across different relative corruption levels with a combined flow-matching and DFM loss.The network predicts denoised translations and rotations and an amino-acid distribution; its architecture modifies FrameFlow with a larger transformer, smaller Invariant Point Attention, and an amino-acid prediction head.Source: §4.2 · p. 6

4.3. Sampling

  • Decoupled sampling schedules support conditional inpainting by fixing either modality; fixing structure generates sequence, while purity selects amino-acid positions to unmask.The paper summarizes flexible multimodal sampling in Figure 1C and Table 2.Source: §4.3 · p. 6, Table 2 · p. 6

5. Related Work

The related work positions DFMs as a flexible continuous-time framework for discrete data and Multiflow as joint protein structure–sequence generation, addressing limitations of prior approaches that handled modalities separately.

Discrete Diffusion Models

  • DFMs extend continuous-time discrete diffusion with directly specified noising distributions and inference-time rate matrices, while encompassing prior discrete diffusion models.Unlike matrix-exponential-restricted processes, DFMs allow choosing the inference-time rate matrix. Igashov et al. (2023) used a data-conditional process without a modeling framework, Zhang et al. (2023) supplied building blocks, and Peng et al. (2023), Vignac et al. (2023b), and Hua et al. (2023) studied multimodal molecular diffusion.Source: §5 · p. 7

Protein Generation

  • Prior diffusion and flow models generated protein backbones, while separate diffusion work targeted sequence generation; the backbone methods required a separate sequence model.RFDiffusion generated proteins validated in wet-lab experiments; the cited backbone-generation work includes Yim et al. (2023b;a), Bose et al. (2023), Lin & AlQuraishi (2023), and Ingraham et al. (2023). Sequence-generation work includes Alamdari et al. (2023), Gruver et al. (2023), Yang et al. (2023), and Yi et al. (2023).Source: §5 · p. 7
  • ProteinGenerator diffuses amino-acid one-hots while predicting structure, whereas Protpardelle diffuses structure while predicting sequence; Multiflow models both jointly for flexible inference conditioning.Luo et al. (2022) and Shi et al. (2022) limited co-design to antibody CDR loops; Anand & Achim (2022) presented diffusion on structure and sequence but reported neither standard evaluation metrics nor available code.Source: §5 · p. 7

6. Experiments

The experiments show that sample-time stochasticity improves discrete text generation and that Multiflow supports joint protein structure–sequence generation. They also examine protein-generation quality, crossmodal control, and folding tasks the same model can perform without retraining.

6.1. Text Modeling

  • For text generation, DFM outperforms D3PM and SEDD because sample-time flexibility lets it optimize the NLL–entropy frontier by choosing η.The reported choice is η = 15; at η = 0, DFM performs similarly to D3PM.Source: §6.1 · p. 7, Figure 2 · p. 7

6.2. Protein generation

  • Protein designability uses scRMSD < 2 Å; Co-design 1 evaluates generated structure–sequence pairs, while PMPNN 1/8 evaluate structures, with diversity and novelty as further checks.Diversity is the number of unique FoldSeek clusters; novelty is the average TM-score to the most similar PDB protein.Source: §6.2 · p. 8
  • Distillation improves overall designability and diversity by replacing training sequences with the lowest-scRMSD PMPNN variants and adding 4,179 accepted synthetic examples.The synthetic examples were generated by Multiflow and passed PMPNN 8 designability.Source: §6.2.1 · p. 8, §6.2 · p. 8

6.2.1. CO-DESIGN RESULTS.

  • Varying sequence CTMC stochasticity changes sampled proteins’ secondary-structure composition, letting the multimodal model tune structural properties at inference time.Figure 3 tracks the proportions of residues in alpha helices and beta strands across stochasticity levels.Source: §6.2.1 · pp. 8–9, Figure 3 · p. 9
  • Multiflow surpasses previous co-design methods and all baselines in PMPNN 8 designability, while its co-designed sequences achieve designability comparable to PMPNN 1.The comparison averages three seeds; Multiflow’s joint generation does not use PMPNN at inference.Source: §6.2.1 · p. 8, Table 3 · p. 9

6.2.2. FORWARD AND INVERSE FOLDING

  • Without retraining, Multiflow performs close to ProteinMPNN but poorly compared with ESMFold on forward and inverse folding.The paper identifies this as a limitation: Multiflow does not perform competitively on every generation task.Source: Table 4 · p. 9, §6.2.2 · p. 9

7. Discussion

Discrete Flow Models (DFMs) provide a flow-based framework for discrete generation that enables multimodal modeling, demonstrated through protein co-design with Multiflow. Future work includes developing domain-specific DFMs and improving Multiflow across protein-generation tasks.

  • DFMs offer a simple-to-implement flow-based framework that removes limitations in defining corruption processes and provides more sampling flexibility for improved performance over previous discrete diffusion models.Source: §7 · p. 9
  • Combining a DFM with FrameFlow enables state-of-the-art protein co-design with Multiflow.Source: §7 · p. 9
  • Future work includes developing more domain-specific DFM models and improving Multiflow on protein-generation tasks, including sidechain modeling.Source: §7 · p. 9

Impact statement … B.1. Proof of Proposition 3.1

The paper flags both beneficial and harmful uses of improved protein modeling and describes safeguards against misuse. Its appendix documents the scope of its supporting analyses and establishes that averaging conditional CTMC rates generates the desired marginals.

Impact statement

  • The authors warn that generative modeling could aid enzyme and drug-candidate design or be misused to design toxic substances, and omit methods easily adapted to harmful design without expert knowledge.They state that care is needed to ensure the broad societal impacts of improved protein modeling are positive.Source: Impact statement · p. 10

A. Organization of Appendix

  • Appendix C analyzes DFM’s cross-entropy objective and its link to model log-likelihood, while Appendices D–F cover related work, multidimensional factorization, and PyTorch realizations.The PyTorch examples use the masking or uniform forms for pt|1.Source: A. Organization of Appendix · p. 14
  • Appendices G–J discuss CTMC sampling, compare DFM with discrete diffusion, and provide further text and protein co-design experiment details and results.Appendix H finds that classical discrete diffusion models in discrete and continuous time fit within the DFM framework.Source: A. Organization of Appendix · p. 14

B. Proofs

B.1. Proof of Proposition 3.1

  • The Proposition 3.1 proof takes pdata expectations of the conditional Kolmogorov equation and uses the diagonal-rate identity Rt(i, i) = −Σj≠i Rt(i, j).The displayed derivation expresses the resulting term using pdata(x1)pt|1(j|x1)Rt(j, xt|x1).Source: §B.1 · p. 14
  • Averaging the conditional rates Ep1|t(x1|xt)[Rt(xt, j|x1)] generates the marginals pt(xt).The proof identifies the resulting expression as the Kolmogorov equation for a CTMC with marginals pt(xt).Source: §B.1 · p. 14

B.2. Proof of Proposition 3.2 … B.6. Proof of Proposition 4.2

These proofs verify conditional and joint flow-generating rates, establish minimum-jump properties for masking and uniform interpolants under factorization assumptions, and delimit when those guarantees apply.

B.2. Proof of Proposition 3.2

  • Proposition 3.2’s conditional rate satisfies the Kolmogorov equation when zero-mass states cannot regain mass; endpoint exceptions for the masking and uniform interpolants can be approximated by starting and stopping at ε offsets.The approximation becomes arbitrarily accurate as ε → 0.Source: §B.2 · p. 15

B.3. Proof of Proposition 3.3

  • Adding any nonnegative multiple of a detailed-balance rate to a rate generating the conditional flow preserves that flow; R*t is a specific case.A detailed-balance rate contributes zero to the conditional time derivative.Source: §B.3 · p. 16

B.4. Proof of Proposition 3.4

  • The optimality results for multidimensional data assume a factorized conditional flow and a rate matrix that permits only single-dimension jumps.Source: §B.4 · p. 16

B.4.1. MASKING INTERPOLANT

  • Under the factorized masking interpolant, R*t makes exactly one jump per dimension and achieves the minimum expected total of D jumps.Source: §B.4.1 · p. 17

B.4.2. UNIFORM INTERPOLANT

  • For the uniform interpolant, R*t makes one jump in each differing coordinate, so its expected jump count equals the expected Hamming distance and is minimal.Source: §B.4.2 · p. 17

B.4.3. DISCUSSION

  • The optimality guarantee holds conditionally only for the two simple flows, not necessarily for unconditional marginals, whose optimal rate may require a global optimization.Source: §B.4.3 · p. 17

B.5. Proof of Proposition 4.1

  • The selected independent velocities and rate matrix for continuous and discrete modalities generate the desired conditional joint flow under the Fokker-Planck-Kolmogorov equation.Source: §B.5 · pp. 18–19

B.6. Proof of Proposition 4.2

  • Taking the expectation of the conditional joint equation and substituting the specified velocities and rate matrix shows that they generate the desired unconditional joint flow.Source: §B.6 · pp. 19–21

C. Analysis of Training Objective … C.2.1. OBJECTIVE FOR THE MASKING INTERPOLANT

The analysis derives a CTMC evidence lower bound (ELBO), decomposes it into cross-entropy, rate-regularization, and KL terms, and shows that for the masking interpolant the ELBO is exactly a weighted cross-entropy. It argues that the rate-regularizing term depends on an arbitrary conditional rate choice and gives the masking loss’s weight and scope.

C. Analysis of Training Objective

  • The analysis examines how the cross-entropy objective relates to data log-likelihood and to the ELBO used to train classical discrete diffusion models.Its proof proceeds from CTMC path measures to an ELBO decomposition and then to the masking-interpolant result.Source: C. Analysis of Training Objective · p. 21

C.1. Introduction to CTMC path measures

  • CTMC path measures assign probabilities to trajectories described by jump times and intervening states, with paths right-continuous and possessing left limits.At a jump time, the path takes the new state while its left limit is the state before the jump.Source: §C.1 · p. 21

C.1.1. DERIVATION OF LELBO

  • The ELBO derivation treats the CTMC’s intermediate states as latent variables and uses an x1-conditioned CTMC with Girsanov’s transformation to obtain the bound.The generative process uses rate matrix Rθ, while the conditioned process begins from p0|1(x0|x1) and uses an x1-conditioned rate matrix.Source: §C.1.1 · p. 23

C.2. Decomposition of LELBO 

  • The model log-likelihood is bounded below by Lce + LR + LKL + C, with LKL enabling cross-entropy to replace the usual ELBO denoising signal and LR acting as a rate-forcing term.The analysis argues that omitting LKL makes the bound looser and that LR encourages the learned rate to match an arbitrarily chosen conditional rate.Source: §C.2 · pp. 25, 28
  • The likelihood-based cross-entropy is weighted by ω(xt), but setting ω(xt) = 1 equally weights loss levels and makes the loss independent of the arbitrary conditional rate matrix.The paper motivates this as an approximation analogous to one used in diffusion-model training.Source: §C.2 · p. 26
  • The rate-forcing term is undesirable because it encourages the learned generative rate to match a jump rate determined by the arbitrary stochasticity of the chosen conditional rate matrix.The paper expects the true LR term to have a similar effect to its approximation, ˆLR.Source: §C.2 · p. 27

C.2.1. OBJECTIVE FOR THE MASKING INTERPOLANT

  • For the masking interpolant, LELBO is exactly cross-entropy weighted by 1/(1-t) and evaluated only on dimensions masked in the corrupted sample.The result assumes multidimensional data and the specified masking conditional flow.Source: §C.2.1 · pp. 28–29

D. Discussion of Related Work

The section positions Discrete Flow Models (DFMs) as a way to bring flow-modeling benefits to discrete data while relaxing constraints on how discrete processes are constructed and sampled. It also describes how learning across modality noise levels enables flexible multimodal generation.

D. Discussion of Related Work

  • DFM extends interpolant-based flow modeling to discrete data by specifying p_t|1 directly and differentiating it for a rate matrix, avoiding matrix-exponential tractability constraints.Prior discrete diffusion methods define noise through a rate matrix and require the matrix exponential to obtain the total noise added from time 1 to t.Source: D. Discussion of Related Work · pp. 29–30
  • DFM permits infinitely many valid sampling rate matrices, making the inference-time process selectable rather than fixing diffusion to its canonical time reversal.The standard discrete-diffusion ELBO depends on the initial rate matrix and assumes sampling by time reversal; prior corrector steps do not remove that canonical-process assumption.Source: D. Discussion of Related Work · p. 30
  • Prior discrete alternatives include categorical score matching, cross-entropy, and data-ratio losses, with differing architecture requirements and likelihood or ELBO properties.Sun et al. (2023b) require specialized architectures; Vignac et al. (2023a) do not analyze the cross-entropy/log-likelihood link; Meng et al. (2022) use an L2 loss that Lou et al. (2023) describe as not penalizing mode dropping, while score entropy requires forward-rate weighting to be a true ELBO.Source: D. Discussion of Related Work · p. 30
  • Prior multimodal diffusion combines continuous and discrete processes for tabular and molecule generation; DFM learns across relative modality noise levels to support co-generation, inverse folding, or forward folding.Prior molecule work also found that corruption order can matter: Peng et al. (2023) corrupt bonds before atom positions, while Vignac et al. (2023b) corrupt atom types, then bonds, then positions.Source: D. Discussion of Related Work · p. 30
  • Embedding discrete variables into continuous diffusion loses their discrete structure, which matters for algorithmically meaningful variables such as whether a graph edge exists.The passage cites Li et al. (2022), Chen et al. (2023a), Richemond et al. (2022), Gong et al. (2023), Dieleman et al. (2022), Han et al. (2022), Strudel et al. (2022), Gulrajani & Hashimoto (2023), and Floto et al. (2023) as examples of embedding approaches.Source: D. Discussion of Related Work · p. 30
  • Chen et al. (2023b) construct discrete diffusion from marginals for masking and uniform cases; DFM can use this approach at η = 0 but also allows η > 0.In the masking case, conditioning on the transition times can reduce sampling to D function evaluations; the passage notes this does not hold for unconditional uniform corruption.Source: D. Discussion of Related Work · p. 31
  • Unlike GFlowNets, which sample from an energy function, DFM targets novel datapoints from training examples and uses detailed balance to increase inference-time flexibility.The passage notes that GFlowNets can also use detailed balance, which is also used in MCMC methods.Source: D. Discussion of Related Work · p. 31

E. Multidimensional Data … E.3. Detailed Balance

The appendix extends DFM to multidimensional categorical data by deriving factorized rates for generative modeling, then describes dimensionwise training and sampling and establishes when detailed balance holds for the full process.

E. Multidimensional Data

  • The derivation assumes independent corruption in each dimension and a rate matrix that sums dimensionwise rates, ruling out transitions that change multiple dimensions at once.The paper notes that simultaneous multidimensional changes are enabled later in the sampling algorithm.Source: E. Multidimensional Data · p. 31
  • The derived multidimensional rate satisfies the unconditional Kolmogorov equation and can be used for generative modeling.The conditional rate is first verified to generate the desired conditional flow, then converted to an unconditional rate.Source: E. Multidimensional Data · pp. 33–34

E.1. Training

  • The training network outputs D × S logits from x_t, converts them with softmax into per-dimension denoising probabilities, and learns them with per-dimension cross-entropy.Source: §E.1 · p. 34

E.2. Sampling

  • The standard Euler transition step cannot change more than one dimension at a time because the learned rate matrix is factorized.Source: §E.2 · p. 34
  • A factorized approximation to the Euler step permits simultaneous changes across dimensions and approximates the standard transition probability to first order.It is equivalent to taking an Euler step independently in each dimension at every simulation step.Source: §E.2 · p. 34

E.3. Detailed Balance

  • If each dimension’s rate matrix is in detailed balance with its conditional marginal, the full-dimensional rate matrix is in detailed balance with the full conditional marginals.When states differ in exactly one dimension, the multidimensional balance equation reduces to the standard one-dimensional equation; when they differ in more than one, both sides are zero.Source: §E.3 · p. 35

F. Implementation Details … F.3.1. DETAILED BALANCE

This section presents implementation details for DFM training and sampling with masking, uniform, and general conditional interpolants, including detailed-balance variants and sampling procedures that can use model predictions.

F. Implementation Details

  • The appendix gives concrete DFM derivations and minimal PyTorch training and sampling loops for masking, uniform, and general conditional paths under multidimensional factorization assumptions.It also points to notebooks containing the minimal examples.Source: F. Implementation Details · p. 35

F.1. Masking Example

  • For masking, the denoiser predicts valid-token logits and is trained on masked dimensions; sampling unmasks with probability dt/(1-t) and leaves already-unmasked values unchanged.The training loop excludes unmasked dimensions from the loss; at sampling, an unmasked position stays in its current state.Source: §F.1 · pp. 36–37

F.1.1. DETAILED BALANCE

  • Detailed-balance masking adds reverse transitions at rate η and corresponding mask-to-token transitions at rate ηt/(1-t), creating stochastic mask/unmask flux while preserving the conditional flow for any η.When a position is masked, its unmasking probability gains a factor 1+ηt/(1-t); an unmasked position can return to the mask at probability ηΔt.Source: §F.1.1 · pp. 38–39

F.1.2. PURITY SAMPLING

  • Purity sampling ranks masked dimensions by model probability and unmasks the highest-ranked dimensions first, with the number selected drawn from a binomial distribution.For η > 0, the binomial success probability becomes Δt(1+ηt)/(1-t), while randomly selected dimensions are remasked with probability Δη.Source: §F.1.2 · pp. 39–40

F.2. Uniform Example

  • For the uniform interpolant, the appendix derives DFM and provides PyTorch training and sampling implementations; the training loop replaces selected dimensions with uniform noise at probability 1-t.Source: §F.2 · pp. 40–41
  • In the uniform example, off-diagonal step probabilities are computed from x1 probabilities and dt/(1−t), then diagonal probabilities are set so each row sums to one.The same row-sum construction is shown for the uniform interpolant with noise.Source: §F.2 · p. 42, §F.2.1 · p. 44

F.2.1. DETAILED BALANCE

  • For the uniform interpolant, detailed balance permits tunable stochasticity by setting a_t = η and determining b_t from the balance equation, producing flux on and off the predicted clean state.The assumed rates include transitions away from the clean state and compensating transitions toward it; η controls the rate of leaving clean data.Source: §F.2.1 · pp. 42–43

F.3. General Case

  • For general conditional flows, training requires a factorized, differentiable interpolant with zero-probability states having zero time derivative, and uses cross-entropy on sampled interpolants.The loop samples data and t, draws x_t from p_t|1(x_t|x_1), and trains predictions against x_1.Source: §F.3 · pp. 44–45
  • General-case sampling computes R*_t from the conditional interpolant and may substitute a sample from the model prediction when its expectation is intractable.The procedure assumes a prior sampler for the chosen interpolant and constructs step probabilities from the resulting transition rates.Source: §F.3 · p. 45

F.3.1. DETAILED BALANCE

  • For general interpolants, detailed-balance solutions depend on the desired noise; one baseline sets the upper-triangular rates to 1 and the lower-triangular rates to their balance ratio.The supplied passage introduces this as one possible choice among many.Source: §F.3.1 · p. 47

G. CTMC Sampling Methods … H.2. Discrete Time Discrete Diffusion Models

The appendices describe DFM sampling and compare its rate and time-discretization choices with continuous- and discrete-time diffusion models.

G. CTMC Sampling Methods

  • Sampling x1 from the learned conditional and then sampling the next state from the joint preserves the same marginal distribution over xt+∆t as the full expectation update.This lets the Euler update avoid calculating the full expectation over conditional rates.Source: G. CTMC Sampling Methods · p. 47
  • τ-leaping permits multiple jumps across dimensions and can cover larger distances for ordinal data, but the update requires the full unconditional rate matrix.Multiple jumps within one dimension are not meaningful for categorical data.Source: G. CTMC Sampling Methods · p. 47
  • Faster CTMC samplers exist, but this work prioritizes framework simplicity over sampling speed and leaves applying those methods to future work.The passage cites Sun et al. (2023b) and Lou et al. (2023) as examples of potentially faster samplers.Source: G. CTMC Sampling Methods · p. 47

H. Comparison with Discrete Diffusion Models

H.1. Continuous Time Discrete Diffusion Models

  • DFM can specify its noise schedule directly and choose the conditional rate matrix at test time, whereas classical discrete diffusion implicitly fixes that rate choice during training.Classical discrete diffusion also faces a harder schedule derivation because it must apply the matrix exponential to its rate matrix.Source: §H.1 · p. 49

H.2. Discrete Time Discrete Diffusion Models

  • DFM permits any time discretization at test time because training covers all t ∈[0, 1], unlike the original D3PM discretization (Austin et al., 2021).The comparison uses the masking process and flips the time convention relative to D3PM, with t = 0 clean and t = T noise.Source: §H.2 · pp. 49–50
  • For the masking process, DFM derives a conditional rate matrix absent from prior D3PM formulations, while both frameworks’ ELBO simplifies to a weighted cross-entropy term (Austin et al., 2021).The passage points to Appendix C.2.1 for the corresponding DFM result.Source: §H.2 · p. 50

I. Text Experiment Details … J.1. Experimental Details

The text experiments show that DFM sample quality depends on the sampling noise level and that its test-set BPC is similar to prior masking-style diffusion models. The appendices also describe the text evaluation protocol and Multiflow protein-generation experiments.

I. Text Experiment Details

  • Text models use a smallest-GPT2-style transformer with 768-dimensional embeddings, 12 blocks and heads, and 28 output categories for lowercase letters, whitespace, and mask.The model uses shared token embeddings for noisy inputs and self-conditioning predictions, learned position embeddings, sinusoidal time embeddings, and QK-layernorm.Source: §I · pp. 50–51
  • Training uses text8 sequences of length 256, batch size 256 with eight accumulation steps, and self-conditioning on every iteration; D3PM uses equal cross-entropy weighting across time.Self-conditioning uses prior predictions on half of iterations and all-mask priors on the other half.Source: §I · p. 51
  • DFM and D3PM train for 750k iterations; autoregressive models are evaluated at minimum validation loss after overfitting begins, while SEDD uses 750k iterations and effective batch size 2048.DFM uses a linear masking interpolant, D3PM an absorbing-state process, and SEDD otherwise retains its default parameters.Source: §I · p. 51
  • For evaluation, DFM uses Δt = 0.001 through t = 0.98, replaces remaining masks with the most likely denoised token, and scores 512 sequences using GPT-J-6B negative log-likelihood and sample entropy.The 512 sequences are each 256 tokens long. D3PM uses 1000 training timesteps to match DFM, and the dataset’s reference negative log-likelihood is 4.2 under GPT-J-6B.Source: §I · pp. 51–52

I.2. Model Log-Likelihoods

  • DFM achieves BPC similar to earlier masking-style diffusion models, while Lou et al. (2023) achieves the lowest BPC.The comparison reports test-set bits-per-character; the supplied passages do not give the table’s individual values.Source: Table 5 · p. 52, §I.2 · p. 52

I.1. Stochasticity Sweep

  • At η = 15, DFM provides good sample quality, whereas η = 50 can produce incoherent text, especially at high logit temperatures.Increasing noise generally lowers negative log-likelihood, but excessive noise can increase it and reduce sample entropy away from the dataset.Source: §I.1 · p. 52, Figure 4 · p. 52

I.3. Example Text Generations

  • For example generations, the authors select each model’s temperature to match sample entropy while allowing negative log-likelihood to vary; SEDD uses a corruption style instead.The selected temperatures are shown in Figure 5.Source: §I.3 · p. 53, Figure 5 · p. 53

J. Protein Generation Experiment Details / J.1. Experimental Details

  • Multiflow protein experiments use a modified FrameDiff architecture with eight blocks, four transformer layers per block, 16 IPA hidden channels, and a three-layer sequence-logit MLP.The model removes skip connections and psi-angle prediction; its corruption-level objective sets one modality’s time to 1 in 10% of cases for each modality, and independently samples both times in the remaining 80%.Source: §J.1 · pp. 54–55

J.2. Additional Multiflow Results … J.4.2. INVERSE FOLDING RESULTS

Multiflow is evaluated as a flexible model for co-generation, forward folding, and inverse folding, with results examining oracle choice, a conditional-flow ablation, and folding performance. The experiments show competitive inverse-folding performance, while forward folding captures some structural properties but often misses exact folds.

J.2. Additional Multiflow Results

  • With AlphaFold2 as the evaluation oracle, designability is lower overall, but distillation yields a two-fold improvement with either oracle, indicating no overfitting to ESMFold.The authors use ESMFold to filter training examples and for their main evaluation; Table 7 reports the Co-design 1 designability results by oracle.Source: Table 7 · p. 56, §J.2 · p. 56

J.3. Uniform Conditional Flow Ablation

  • The uniform conditional flow has a worse trade-off between Co-Design 1 designability and diversity than the masking interpolant used in the main models.With greater stochasticity, designability generally rises while structural-cluster diversity falls; high stochasticity favors alpha helices over beta strands.Source: §J.3 · p. 56, Figure 8 · p. 56

J.4. Forward and Inverse Folding Experiments

  • Multiflow can co-generate sequence and structure and perform forward or inverse folding by choosing which modality to provide and which to generate.Unlike specialized forward- and inverse-folding models, Multiflow performs all three tasks; improving its folding performance is left to future work.Source: §J.4 · p. 57

J.4.2. INVERSE FOLDING RESULTS

  • For inverse folding, Multiflow receives the ground-truth structure and generates a sequence, assessed by the scRMSD of the structure predicted from that sequence.The objective is to minimize scRMSD, the RMSD between the ESMFold-predicted structure and the input structure.Source: §J.4 · p. 57, §J.4.2 · p. 58

J.4. Forward and Inverse Folding Experiments

  • The folding comparison is an initial baseline because ESMFold and ProteinMPNN use different training and test sets that the authors did not retrain to match Multiflow.The 449-protein test set uses a time-based PDB split and filtering to exclude proteins present in the training data of Multiflow, ProteinMPNN, or ESMFold.Source: §J.4 · p. 57

J.4.1. FORWARD FOLDING RESULTS

  • In forward folding, Multiflow has some success on shorter proteins but struggles on longer ones, often producing realistic structures without accurately recovering the exact fold.Forward folding supplies the ground-truth sequence and runs the structure flow; predicted and ground-truth secondary-structure proportions correlate.Source: §J.4.1 · pp. 57–58, Figure 9 · p. 58, Figure 10 · p. 58

J.4.2. INVERSE FOLDING RESULTS

  • Multiflow performs competitively with ProteinMPNN across protein lengths, although ProteinMPNN achieves slightly lower scRMSD on average.Both models’ scRMSD values tend to cluster around 1 to 2; each has proteins on which it achieves the lower value.Source: §J.4.2 · p. 58, Figure 11 · p. 59
Loading 2402.04997v2…