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

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.

한 줄 요약

Discrete Flow Models(DFMs)를 제안하여, 연속과 이산 데이터를 결합한 다중모달 생성 모델을 구축하고 단백질 공설계에 적용한다.

핵심 기여도

Discrete Flow Models(DFMs)를 제안하여, 이산 데이터에 대한 flow-based 모델링을 CTMC 기반으로 구현.
기존 이산 확산 모델을 포함하면서도 더 높은 성능과 샘플링 유연성을 제공.
Multiflow라는 다중모달 생성 모델을 개발하여, 단백질 구조와 서열을 동시에 생성.
CTMC Stochasticity 조절로 샘플의 이차 구조 구성 및 다양성을 제어 가능.

핵심 아이디어

DFMs는 이산 공간에서의 flow matching을 CTMC(Continuous Time Markov Chains)를 통해 구현하는 새로운 이산 생성 모델이다. 기존 확산 모델은 샘플링 시간 유연성이 낮고, 모델 재학습이 필요하지만, DFMs는 샘플링 시간을 조절하면서도 훈련 없이 성능을 유지할 수 있다. 이는 확률 흐름 $ p_t $를 정의하고, 이 흐름을 따라 $ x_t $를 시뮬레이션하는 방식으로 이루어진다. CTMC Stochasticity는 샘플링 경로의 변이성을 조절하여 샘플의 분포적 특성을 제어할 수 있다. 이는 기존 이산 확산 모델(D3PM)보다 샘플링 유연성과 샘플 품질을 개선하는 핵심 아이디어이다.

기술적 접근법

**DFMs**: CTMC 기반의 이산 확률 흐름 모델.
**Multiflow**: DFM과 FrameFlow를 결합한 다중모달 생성 모델.
**샘플링 유연성**: CTMC Stochasticity 조절로 샘플링 경로의 변이성 조절.
**네트워크 아키텍처**: FrameDiff 기반으로 수정된 모델. IPA(Invariant Point Attention)와 트랜스포머 블록 사용.
**훈련 전략**: $ t, \tilde{t} $를 다양한 수준의 오염 비율로 샘플링하여 모델이 구조와 서열 간 상호작용을 학습하도록 유도.

주요 결과

**Text 데이터 실험**: DFMs는 D3PM 대비 샘플링 유연성 확대로 샘플링 성능 향상.
**단백질 공설계**: Multiflow는 기존 방법 대비 최고 성능 달성.
**구조 생성**: 데이터 디스틸레이션을 통해 최고 수준의 구조 생성 성능 달성.
**CTMC Stochasticity**: 샘플의 이차 구조 구성 및 다양성 제어 가능.

의의 및 한계

DFMs는 이산 데이터 생성 모델의 샘플링 유연성을 획기적으로 향상시키며, 다중모달 생성 문제를 해결하는 새로운 기반을 제공한다. 특히, Multiflow는 단백질 구조와 서열을 동시에 생성할 수 있어, 단백질 공설계 분야에서 중요한 진전이다. 그러나, 본 연구는 sidechain 모델링 등 일부 단백질 생성 작업에서의 성능 개선이 필요하다는 한계를 인정하고 있다. 또한, CTMC 기반의 흐름 모델링은 복잡한 이산 공간에서의 확률적 경로 추정에 기반하기 때문에, 계산 효율성과 모델 해석 가능성에 대한 추가 연구가 필요하다.

실용적 활용

DFMs는 텍스트, 단백질 서열 등 이산 데이터 생성에 활용 가능하며, Multiflow는 단백질 공설계, 의약품 개발, 바이오엔지니어링 분야에서 유용하게 사용될 수 있다. 특히, CTMC Stochasticity 조절 기능은 샘플 다양성 제어가 필요한 생성 작업에 적합하다.