[논문 리뷰] Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design

Andrew T. Campbell, Jason Yim외 3명|PubMed|2024. 02. 07.

Cancer-related gene regulation인용 수 9

한 줄 요약

Discret Flow Models (DFMs)를 도입하여 이산 데이터와 연속 데이터를 모두 다루는 멀티모달 generative 모델링을 가능하게 하고; 구조-시퀀스 단일 프레임워크에서 단독으로 최첨단 성과를 달성하는 Multiflow를 제시합니다.

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.

연구 동기 및 목표

• 이산 데이터와 연속 데이터를 하나의 프레임워크로 통합하여 멀티모달 생성 모델링의 동기를 부여합니다.
• 연속-시간 마르코프 체인을 기반으로 DFMs를 개발하여 이산 데이터 흐름을 모델링합니다.
• DFMs와 연속 상태 흐름 모델을 결합하여 멀티모달 프레임워크를 창출합니다.
• 프레임워크를 단백질 공동 설계에 적용하여 단백질 구조와 서열을 함께 생성합니다.
• 멀티모달 생성에서 최첨단 성능과 샘플링의 유연성을 입증합니다.

제안 방법

• 노이즈에서 데이터로 보간되도록 확률 흐름 p_t를 데이터 조건부 흐름 p_{t|1}(·|x_1)을 통해 정의합니다.
• p_t를 조건부 흐름의 기대로 표현하여 닫힌 형태의 구성과 계산 가능한 학습을 가능하게 합니다(Eq. 6).
• p_{1|t}의 기대값이 p_t를 생성하도록 p_{1|t}에 대해 기대하는 속도 행렬 R_t(x_t, j|x_1)을 구성합니다(Prop. 3.1).
• 교차 엔트로피를 사용하여 디노이즈 분포를 근사하는 신경 네거이저 p_{1|t}^θ(x_1|x_t)를 학습합니다(Eq. 10).
• 시작 속도 R_t^*를 도입하고 자세-수렴 구성 요소 R_t^{DB}를 확장하여 CTMC의 확률적 성격을 제어하는 가족 R_t^η를 형성합니다(Eq. 14).
• DFMs를 연속 흐름 구조 모델과 결합하여 단일 구조-시퀀스 생성을 위한 다중 모달 흐름(Multiflow)을 구성합니다(섹션 4).
• 특정 모달리티를 단백질에 적용하기 위해 서열은 DFM으로 모델링하고 구조는 FrameFlow와 유사한 연속 흐름으로 구성하여 공동 설계를 가능하게 합니다(섹션 4).

실험 결과

연구 질문

• RQ1 연속-시간 마르코프 체인으로 구성된 이산 흐름 모델이 연속 흐름 모델의 유연성으로 이산 데이터를 샘플링할 수 있게 할 수 있을까?
• RQ2 DFMs를 연속 흐름 모델과 결합하면 멀티모달 작업에서 효과적인 멀티모달 생성 및 샘플링 제어가 가능할까?
• RQ3 제안된 Multiflow 프레임워크가 구조와 서열을 별도로 모델링한 이전 방법들에 비해 단백질 공동 설계에서 얼마나 성능이 우수한가?
• RQ4 CTMC의 확률적 특성(η에 따른 영향)이 단백질 공동 설계에서 샘플 품질, 다양성 및 설계 가능성에 어떤 영향을 미치는가?
• RQ5 프레임워크가 모달리티 간 조건화 및 인페인팅(예: 서열을 주면 구조를 생성하거나 그 반대)을 지원할 수 있는가?

주요 결과

• DFMs는 CTMC 확률적 특성을 통해 이산 데이터 생성을 샘플링 시간에 따라 유연하게 제어할 수 있으며, η에 적합한 경우 텍스트 데이터에서 이산 확산 기반보다 더 우수한 성능을 보입니다.
• Multiflow는 단백질 구조와 서열을 함께 생성하는 공동 설계에서 최첨단 성능을 달성합니다.
• 모델은 어느 모달리티에 대해서도 조건화 및 유연한 멀티모달 샘플링(인페인팅 및 순방향/역방향 접힘 작업 포함)을 지원합니다.
• CTMC 확률성(η로 제어)은 샘플 다양성과 점프 빈도에 영향을 주어 탐색과 충실도 간의 균형을 조정하는 조정기 역할을 합니다.
• 예비 결과는 Multiflow가 역/전향 접힘에도 잠재력을 지니며 일반 용도의 단백질 생성 모델링으로의 경로를 시사합니다.