PaperlayerView exact revision on arXiv

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

The paper addresses the challenge of generative modeling across discrete and continuous data, especially for jointly designing protein structures and sequences. It introduces CTMC-based Discrete Flow Models and combines them with continuous flows in Multiflow. Multiflow achieves state-of-the-art protein co-design performance while supporting flexible generation and conditioning across modalities.

  • Problem

    Generative models need to handle discrete and continuous data jointly, but diffusion-based multimodal modeling has inflexible sampling requirements that can require retraining and extensive evaluation.

  • Method

    Discrete Flow Models use Continuous Time Markov Chains to model discrete probability flows and combine with continuous flows through factorized multimodal modeling.

  • Results

    Multiflow achieves state-of-the-art co-design performance for jointly generating protein structure and sequence, with sampling flexibility that can control sample properties.

  • Takeaways & Limitations

    The framework supports one multimodal protein model for joint generation and conditioning on either structure or sequence, with preliminary forward- and inverse-folding results.

  • Takeaways & Limitations

    The paper does not present specific Multiflow methods for harmful-substance design without expert knowledge and identifies improving performance across protein-generation tasks as future work.

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 paper introduces Discrete Flow Models (DFMs) to extend flow-based generation to discrete data and combines them with continuous flows for multimodal protein co-design. Multiflow jointly generates protein structures and sequences while supporting conditioning on either modality and flexible sampling.

  • Motivation: Multimodal generative models are needed to jointly model continuous protein structures and discrete amino acid sequences for protein co-design.Protein structure endows function, while sequence provides the blueprint for how structure is made.
  • Motivation: Diffusion models support continuous and discrete spaces but require inflexible sampling choices and repeated retraining or evaluation to find suitable parameters.The challenge is amplified when multiple modalities are modeled together.
  • Discrete Flow Models: Discrete Flow Models realize discrete flow-based modeling with Continuous Time Markov Chains, providing sampling flexibility without retraining.DFMs use a probability flow and a denoising neural network trained with cross-entropy; CTMC stochasticity can vary the number of trajectory transitions.
  • Multiflow: Multiflow jointly generates protein structure and sequence while conditioning on either modality, supporting structure or sequence generation, forward or inverse folding, and co-generation.A single model can start sampling from noise or from a provided structure or sequence.
  • Results: Multiflow achieves state-of-the-art co-design performance, while CTMC stochasticity enables control over sample properties such as secondary-structure composition and diversity.Experiments also report improved discrete text modeling over D3PM through expanded sample-time flexibility and preliminary promise for forward and inverse folding.
  • Multimodal framework: The authors combine DFMs with continuous flow methods to create a multimodal generative framework.The framework defines factorized flows for each data modality.

2. Background

Continuous Time Markov Chains model discrete sequences as trajectories that remain in states and jump according to time-dependent rates. Their marginal distributions evolve through a probability-flow ODE.

  • Continuous Time Markov Chains: Over an infinitesimal interval dt, the process either stays in its current state or jumps to another state with probability Rt(xt, j)dt.The transition probabilities form a categorical distribution after including the diagonal stay probability.
  • Continuous Time Markov Chains: Finite-time sequence trajectories can be simulated with Euler steps starting from a sample x0 drawn from p0.The initial distribution and rate matrix together determine the CTMC.
  • Continuous Time Markov Chains: A CTMC is specified by an initial distribution p0 and a rate matrix Rt governing state transitions.Off-diagonal rates are non-negative, while the diagonal rate balances outgoing probability.
  • Kolmogorov equation: The Kolmogorov equation relates the rate matrix to the marginal distribution’s time derivative through incoming minus outgoing probability mass.With probability vectors, this becomes ∂tpt = R⊤t pt.
  • Probability flow: The resulting sequence of marginal distributions pt is called a probability flow.A rate matrix generates the flow when it satisfies the marginal evolution equation.

3. Discrete Flow Models

Discrete Flow Models construct a discrete probability flow from noise to data and sample it with a CTMC. Their conditional-flow formulation permits multiple valid rate matrices and inference-time stochasticity choices.

  • Flow construction: A DFM interpolates between noise at t = 0 and data at t = 1 using a discrete probability flow sampled by a CTMC.The conditional flow is designed to be explicit and converge to the datapoint at t = 1.
  • Flow construction: Conditional flows provide a tractable route to the generative flow because they can be written explicitly when conditioned on a datapoint.The paper uses uniform or mask-state noise in its examples.
  • Rate matrices: The unconditional rate matrix is obtained as an expectation of a conditional rate matrix under the denoising distribution p1|t(x1|xt).A neural network is trained with cross-entropy to approximate this clean-data distribution from noisy inputs.
  • Rate matrices: The cross-entropy objective is independent of the conditional rate matrix, allowing that matrix to be selected after training for inference-time flexibility.Multiple rate matrices can generate the same conditional flow.
  • CTMC stochasticity: Adding a detailed-balance rate matrix preserves the conditional probability flow while changing CTMC trajectory dynamics.The parameter η controls jump frequency and therefore CTMC stochasticity.
  • Rate matrices: The proposed starting rate matrix generates the conditional flow while minimizing the expected number of jumps under the stated factorization assumptions.The construction assumes zero-mass states have zero time derivative and uses non-negative off-diagonal rates.

4. Multimodal Protein Generative Model

The multimodal protein model combines a discrete flow for amino-acid sequences with a continuous flow for structures. Independent noise levels enable flexible sampling and conditional inpainting across modalities.

  • Model representation: Multiflow jointly generates amino-acid sequences and protein structures by combining a DFM with a continuous structure-generation flow.The structure is represented by backbone coordinates and local-frame rotations, while amino acids form the discrete modality.
  • Conditional flow: The conditional multimodal flow factorizes over dimensions and modalities, with continuous ODEs for coordinates and rotations and a CTMC for amino acids.Rotations are modeled on SO(3), while translations use a Euclidean vector field.
  • Conditional flow: Proposition 4.1 verifies that the multimodal process defined by Eq. (15) has the intended conditional flow pt|1(Tt|T1).The proposition establishes consistency between the process construction and the specified multimodal flow.
  • Unconditional flow: Although the conditional flow factorizes, the unconditional generative flow couples modalities and dimensions through predictions based on the full corrupted protein state.Each velocity and rate matrix depends on the entire state Tt.
  • Flexible sampling: Decoupled structure and sequence noise levels allow arbitrary combinations of corruption times during sampling.Setting one modality’s time to 1 supports conditional inpainting while the other modality is updated.
  • Training: Training predicts denoised translations, rotations, and amino-acid distributions from the noised protein, using objectives corresponding to continuous flow matching and the DFM loss.The independent-time objective exposes the model to different relative corruption levels across modalities.

5. Related Work

Prior discrete diffusion methods constrain the corruption or sampling process, while earlier protein co-design systems generated one modality and predicted the other. Multiflow instead models both modalities generatively.

  • Discrete diffusion: DFMs extend continuous-time discrete diffusion by allowing directly specified probability paths and inference-time choices of rate matrices.The framework includes prior discrete diffusion models as a special case.
  • Protein generation: Earlier protein generation methods commonly required separate sequence-generation models when generating protein backbones.RFDiffusion is cited as a successful backbone-generation example with this separation.
  • Protein co-design: ProteinGenerator diffuses amino-acid representations while predicting structure, whereas Protpardelle diffuses structure while iteratively predicting sequence.Both approaches use a generative process for one modality and prediction for the other.
  • Protein co-design: Multiflow uses a generative model over both structure and sequence, allowing flexible conditioning at inference time.This distinguishes it from prior approaches that predict the remaining modality.

6. Experiments

Experiments show that sampling-time stochasticity improves discrete generation, while Multiflow achieves strong protein co-design and supports crossmodal control. Multiflow also approaches specialized inverse-folding performance but is weaker on forward folding.

  • 6.1. Text Modeling: DFM outperforms discrete diffusion baselines in text modeling by exploiting sampling-time flexibility.The comparison includes D3PM and SEDD across NLL and entropy, with temperature sweeps for models that support logit scaling.
  • 6.2.1. Co-design Results: Multiflow jointly generates protein structures and sequences, unlike prior approaches that generate one modality and predict the other.Its evaluation includes co-design, structure-only generation, and crossmodal analyses.
  • 6.2.1. Co-design Results: Multiflow surpasses previous co-design methods and matches PMPNN at a comparable level of sequence–structure designability.The reported results average three seeds, and PMPNN-based metrics distinguish co-design evaluation from structure-only evaluation.
  • 6.2.1. Co-design Results: Multiflow outperforms all baselines on PMPNN 8 structure designability while retaining comparable diversity and novelty.Distillation improves overall designability and diversity relative to the non-distilled model.
  • 6.2.1. Co-design Results: Changing sequence CTMC stochasticity changes sampled proteins’ secondary-structure composition, enabling inference-time crossmodal modulation.Figure 3 reports average alpha-helix and beta-strand proportions with standard-error bars.
  • 6.2.2. Forward and Inverse Folding: Multiflow performs close to ProteinMPNN on inverse folding but poorly relative to ESMFold on forward folding.The same co-design model weights are evaluated without additional training, revealing a limitation across generation tasks.

7. Discussion

The paper presents DFMs as a simple flow-based framework for discrete data and combines them with continuous flows in Multiflow for protein co-design.

  • 7. Discussion: DFMs remove restrictions on corruption processes and provide sampling flexibility for multimodal generation.The framework combines with FrameFlow to produce state-of-the-art protein co-design performance.

Impact statement

The paper advances general-purpose generative modeling for discrete and multimodal data while recognizing potential benefits and misuse risks in protein generation.

  • Impact statement: Improved protein modeling could support enzyme and drug-candidate design but could also be misused to design toxic substances.The authors therefore avoid presenting methods that would make harmful-substance design readily accessible without expert knowledge.
  • Impact statement: The paper limits disclosure of applications that could be easily adapted to harmful-substance design without expert knowledge.This is presented as a mitigation measure for risks from general-purpose protein-generation techniques.

B.2. Proof of Proposition 3.2

The proof verifies that the proposed rate matrix generates the desired conditional flow under a no-resurrection assumption. It then establishes minimum-transition results for masking and uniform interpolants, while identifying limits to the optimality claim.

  • Proof assumptions: The proof assumes that states with zero conditional probability cannot later gain probability mass.This no-resurrection condition ensures the rate construction remains compatible with the Kolmogorov equation.
  • Proof of flow generation: The proposed rate matrix satisfies the conditional Kolmogorov equation and therefore generates the desired conditional flow.The verification covers positive-mass states and the zero-mass case, where both sides of the equation vanish.
  • Transition optimality: For the masking interpolant, the multidimensional process makes exactly one jump per changed dimension, achieving the minimum expected number of transitions.Under factorization, no jump can change more than one dimension, so the minimum is the Hamming distance between endpoints.
  • Transition optimality: The same minimum-transition result is established for the uniform interpolant under the stated factorization construction.Each dimension propagates independently, and dimensions with matching endpoints require no jumps.
  • Scope of the result: The optimality proof applies only to the two simple conditional flows studied and need not hold for arbitrary conditional flows or unconditional rates.General optimal rates may require solving a global programming problem, and unconditional averaging can admit lower-jump alternatives.

C. Analysis of Training Objective

This section derives a CTMC path-space evidence lower bound for the generative model and relates it to the cross-entropy training objective. The derivation removes an arbitrary rate-matching term and yields a practical unweighted objective for masking corruption.

  • ELBO decomposition: The model ELBO decomposes into a cross-entropy term, a rate regularizer, and a KL term.The conditioned CTMC supplies the comparison path measure used in this decomposition.
  • CTMC path measures: The derivation represents CTMC trajectories through jump times and visited states, defining path-space measures for the generative and conditioned processes.These measures support the subsequent change-of-measure and ELBO derivation.
  • Path-space transformation: Girsanov’s transformation changes expectations between CTMC path measures with different rate matrices, enabling the generative-process likelihood analysis.The Radon–Nikodym derivative is expressed as a ratio of path measures.
  • Objective simplification: The rate regularizer is unnecessary because the reference conditional rate matrix is arbitrary and its jump rate depends on the chosen stochasticity.The resulting objective retains the standard cross-entropy training signal rather than forcing the learned rate to match that arbitrary rate.

D. Discussion of Related Work

The paper positions DFMs as a flow-based alternative for discrete data that addresses sampling inflexibility in multimodal diffusion models. DFMs use CTMCs to separate the learned probability flow from the sampling process, while retaining discrete structure and connecting to continuous flows.

  • Limitations of discrete diffusion: Discrete diffusion models constrain sampling through a forward noising process whose matrix exponential must remain tractable.This restricts the choice of forward process and ties training to the selected simulation process.
  • Discrete Flow Models: DFMs define the conditional probability flow directly and derive a compatible rate matrix by differentiating that flow.This avoids requiring the desired noise schedule to arise from a tractable matrix exponential.
  • Sampling flexibility: DFMs permit infinitely many valid sampling rate matrices, allowing the simulation process to be chosen at inference time without retraining.Classical diffusion instead centers on a canonical time-reversal process because its rate matrix enters the ELBO objective.
  • Related approaches: The framework differs from continuous embeddings by preserving discrete structure during generation, which can matter when variables represent inherently discrete quantities.Prior multimodal diffusion methods separately process continuous and discrete features or embed discrete data into continuous spaces.
  • Multimodal modeling: DFMs extend flow-based modeling to discrete states while allowing combination with continuous state-space flows for multimodal generation.The framework is intended for generative modeling from training data, rather than energy-based sampling as in GFlowNets.

E. Multidimensional Data

The multidimensional DFM construction factorizes the discrete process across dimensions under explicit assumptions. It derives per-dimension denoising rates, combines them into a full rate matrix, and uses independent updates or a factorized approximation for simultaneous changes.

  • Assumptions: The multidimensional construction models x1 ∈ {1, . . . , S}D using factorization assumptions applied independently across dimensions.The first assumption creates independent corruption processes, while the third decomposes the rate into per-dimension matrices.
  • Rate factorization: Under the factorized rate assumption, transitions changing more than one dimension simultaneously have zero rate in the standard Euler formulation.The paper later introduces a factorized transition approximation to permit simultaneous multidimensional changes.
  • Denoising model: A neural network outputs D × S logits, which are transformed by softmax into conditionally independent denoising probabilities for each dimension.The network is trained with cross entropy for each dimension.
  • Sampling approximation: The factorized transition approximation performs an Euler step independently in each dimension while matching the full transition distribution to first order.This enables multiple dimensions to change simultaneously in one simulation update.
  • Detailed balance: Detailed balance in every dimension implies detailed balance for the full multidimensional process under the factorized construction.When states differ in one dimension, the multidimensional condition reduces to the corresponding single-dimensional equation.

F. Implementation Details

The implementation derives masking-based DFM training and sampling, then adds detailed-balance noise to vary CTMC stochasticity. The resulting sampler can unmask or remask positions while preserving the target conditional flow.

  • Implementation scope: The derivations assume multidimensional factorization and include masking, uniform, and general conditional-flow implementations with minimal PyTorch examples.The paper notes that the rate construction is valid for t ∈ (0, 1), while exact endpoint cases require separate discussion.
  • Training: Masking DFMs train a denoising network on masked positions using cross entropy over valid unmasked states.The model receives noisy states and time, outputs logits over valid values, and ignores already unmasked dimensions in the loss.
  • Basic sampling: The basic masking sampler repeatedly predicts clean-state probabilities and unmaskes currently masked positions with time-dependent probability.Already unmasked positions remain in their current state.
  • Stochastic sampling: A detailed-balance rate matrix introduces a stochasticity parameter η that allows transitions between masked and unmasked states.Increasing η adds remasking and bonus unmasking rates while maintaining the desired conditional flow.
  • Sampling variants: The factorized implementation can use independent per-dimension updates and a sampling variant that ranks dimensions by model confidence before unmasking.Purity sampling replaces uniform random selection with confidence-based ranking.

G. CTMC Sampling Methods

DFM sampling constructs a generative CTMC from learned denoising distributions and conditional rate matrices. The framework also clarifies that discrete diffusion corresponds to a particular training-time stochasticity choice, whereas DFM permits inference-time discretization and sampling flexibility.

  • DFM sampling: DFM sampling can draw a clean-state estimate and use it in an Euler update instead of computing the full expectation over the conditional rate matrix.The resulting sample has the same distribution as the corresponding marginal of the joint update construction.
  • Alternative CTMC samplers: τ-leaping can make multiple jumps across dimensions, but multiple jumps within one categorical dimension are unsuitable because categorical states have no ordering.The paper notes this restriction does not apply in the same way to ordinal data.
  • Discrete diffusion comparison: Classical discrete diffusion fixes its rate matrix through a corruption process and uses its time reversal as the generative process.The rate matrix must support tractable matrix-exponential calculations for the noise marginals.
  • Sampling flexibility: Classical discrete diffusion corresponds to a particular CTMC stochasticity choice made during training, whereas DFM can select the sampling process at inference.DFM also permits directly specifying the conditional noise schedule without requiring a matrix-exponential-friendly rate matrix.

I. Text Experiment Details

The text experiments evaluate discrete flow models by varying sampling noise and temperature, comparing sample quality with discrete diffusion baselines using likelihood and entropy measures.

  • Training setup: The DFM and D3PM models are trained for 750k iterations with matched optimization settings and 1000 warm-up timesteps.Both use 4 Nvidia A40 GPUs, a learning rate of 10^-4, cosine decay toward 10^-5, AdamW, and weight decay 0.1.
  • Evaluation: Text samples contain 512 sequences of length 256 for each temperature setting applied to the denoising logits.Sample likelihoods are evaluated with GPT-J-6B after BPE tokenization, while entropy is estimated from the empirical token distribution.
  • Sampling analysis: The Entropy-NLL curves vary temperature from 0.5 to 1.0 for each sampling noise level.The selected temperature settings for later generations are highlighted separately for each model.
  • Sampling analysis: Increasing noise generally lowers negative log-likelihood, but excessive noise can produce degenerate samples whose entropy moves away from the dataset distribution.At η = 50, high logit temperatures increase negative log-likelihood and decrease sample entropy.
  • Results: DFM outperforms D3PM and SEDD in the NLL-versus-entropy comparison, while achieving similar bits-per-character to prior masking-based diffusion models.The text comparison attributes the NLL-versus-entropy advantage to additional sample-time flexibility; Lou et al. achieve the lowest reported BPC in the separate likelihood comparison.

J.2. Additional Multiflow Results

Additional Multiflow analyses examine protein-generation metrics, conditioning modes, ablations, and folding behavior across structure and sequence tasks. The model shows competitive or strong performance in several settings, while forward folding remains weaker for longer proteins and specialized-model comparisons are limited by differing training sets.

  • Distillation: Distillation improves designability by filtering training examples and adding Multiflow-generated structures that pass ProteinMPNN and ESMFold criteria.Using AlphaFold2 as an alternative evaluation oracle still gives a two-fold improvement, indicating the gain is not specific to the ESMFold oracle.
  • Uniform conditional flow ablation: The uniform interpolant produces a worse diversity–designability trade-off than the masking interpolant, motivating masking in the main models.The ablation varies CTMC stochasticity and measures both Co-Design 1 metrics and secondary-structure composition.
  • Multimodal generation: Multiflow can perform co-generation, forward folding, and inverse folding, with inference-time flexibility over which protein modality is provided or generated.Forward folding fixes the sequence and samples structure, while inverse folding uses structure to generate a sequence evaluated through self-consistency RMSD.
  • Evaluation scope: Comparisons with ESMFold and ProteinMPNN provide an initial baseline because those specialized models use different training and test sets.Retraining the specialized models on Multiflow’s data was outside the study’s scope.
  • Forward folding results: Multiflow generally struggles with exact folded structures for longer proteins, despite often matching ground-truth secondary-structure distributions.Secondary-structure proportions correlate well with ground truth, but the predicted structures frequently have high RMSD error.
  • Inverse folding results: Multiflow performs competitively with PMPNN across protein lengths for inverse folding, although PMPNN achieves slightly lower scRMSD on average.Both methods commonly produce scRMSD values around 1 to 2, with each outperforming the other on some test proteins.
Loading 2402.04997v2…