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

Andrew T. Campbell, Jason Yim+3 more|PubMed|Feb 7, 2024

Cancer-related gene regulation9 citations

TL;DR

Introduces Discrete Flow Models (DFMs) to enable multimodal generation over discrete and continuous data; presents Multiflow for joint structure-sequence protein design with state-of-the-art results.

ABSTRACT

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.

Motivation & Objective

• Motivate multimodal generative modeling by unifying discrete and continuous data within a single framework.
• Develop Discrete Flow Models (DFMs) based on continuous-time Markov chains to model discrete data flows.
• Create a multimodal framework by combining DFMs with continuous-state flow models.
• Apply the framework to protein co-design to jointly generate protein structure and sequence.
• Demonstrate state-of-the-art performance and sampling flexibility in multimodal generation.

Proposed method

• Define a probability flow p_t that interpolates from noise to data via a data-conditioned flow p_{t|1}(·|x_1).
• Represent p_t as an expectation over conditional flows to enable closed-form construction and tractable training (Eq. 6).
• Instantiate a rate matrix R_t(x_t, j|x_1) whose expectation over p_{1|t} generates p_t (Prop. 3.1).
• Train a neural denoiser p_{1|t}^θ(x_1|x_t) using cross-entropy to approximate the denoising distribution (Eq. 10).
• Introduce a starting rate R_t^* and extend with detailed-balance components R_t^{DB} to form a family R_t^η that controls CTMC stochasticity (Eq. 14).
• Combine DFMs with a continuous-flow structure model to form a Multimodal Flow (Multiflow) for joint structure and sequence generation (Section 4).
• Apply specific modalities to proteins by modeling sequence with DFM and structure with a FrameFlow-like continuous flow, enabling co-design (Section 4).

Experimental results

Research questions

• RQ1 Can a discrete flow model built from continuous-time Markov chains enable sampling of discrete data with the flexibility of continuous-flow models?
• RQ2 Does combining DFMs with continuous-flow models enable effective multimodal generation and sampling control for multimodal tasks?
• RQ3 How does the proposed Multiflow framework perform on protein co-design compared to prior methods that model structure and sequence separately?
• RQ4 To what extent does CTMC stochasticity (via η) influence sample quality, diversity, and designability in protein co-design?
• RQ5 Can the framework support conditioning and inpainting across modalities (e.g., generating structure given sequence or vice versa)?

Key findings

• DFMs enable discrete data generation with sampling-time flexibility via CTMC stochasticity, outperforming discrete diffusion baselines on text data for suitable η.
• Multiflow achieves state-of-the-art co-design performance for jointly generating protein structure and sequence.
• The model supports conditioning on either modality and flexible multimodal sampling, including inpainting and forward/inverse folding tasks.
• CTMC stochasticity (controlled by η) influences sample diversity and jump frequency, providing a knob to balance exploration and fidelity.
• Preliminary results indicate Multiflow’s potential for inverse and forward folding, suggesting a path toward general-purpose protein generative modeling.