[論文レビュー] 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 regulation被引用数 9

ひとこと要約

離散 Flow Models (DFMs) を導入し、離散データと連続データの多模態生成を可能にする；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.

研究の動機と目的

• 離散データと連続データを単一のフレームワーク内で統合することで、多模态生成モデルの動機付けを行う。
• 連続時間マルコフ連鎖に基づいて離散データのフローをモデル化するための Discrete Flow Models (DFMs) を開発する。
• DFMsを連続状態フロー模型と組み合わせて多模态フレームワークを構築する。
• このフレームワークをタンパク質の共設計へ適用し、タンパク質の構造と配列を共同生成する。
• 多模态生成における最先端の性能とサンプリング柔軟性を示す。

提案手法

• ノイズからデータへ補間する確率フロー p_t を、データ条件付きフロー p_{t|1}(·|x_1) を介して定義する。
• p_t を条件付きフローの期待として表現し、閉形式の構成と扱いやすい訓練を可能にする（Eq. 6）。
• p_{1|t} の期待 over で p_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 を連続フロー構造モデルと組み合わせて、構造と配列の共同生成を行う Multimodal Flow（Multiflow）を形成する（セクション4）。
• タンパク質に特定のモダリティを適用し、配列を DFM で、構造を FrameFlow に類する連続フローでモデリングして共設計を可能にする（セクション4）。

実験結果

リサーチクエスチョン

• RQ1 連続時間マルコフ連鎖から構築された離散フローモデルは、連続フローの柔軟性を持って離散データのサンプリングを可能にするか。
• RQ2 DFMs と連続フロー模型を組み合わせることで、多模态タスクの効果的な多模态生成とサンプリング制御を実現できるか。
• RQ3 構造と配列を別々にモデリングする前の方法と比べて、提案された Multiflow フレームワークはタンパク質共設計でどのように性能を示すか。
• RQ4 CTMC の確率性（η を介して）は、タンパク質共設計におけるサンプル品質・多様性・設計可能性にどの程度影響するか。
• RQ5 モダリティ間の条件付けやインペインティング（例：配列から構造を生成、あるいはその逆）をフレームワークがサポートできるか。

主な発見

方法	共設計 1 (Des.)	共設計 1 (Div.)	共設計 1 (Nov.)	PMPNN 8 (Des.)	PMPNN 8 (Div.)	PMPNN 8 (Nov.)	PMPNN 1 (Des.)	PMPNN 1 (Div.)	PMPNN 1 (Nov.)
Protpardelle	0.05	6	0.75	0.92	46	0.67	0.63	33	0.68
ProteinGenerator	0.34	31	0.74	0.88	73	0.71	0.75	56	0.72
RFdiffusion	N/A	N/A	N/A	0.90	161	0.69	0.69	120	0.70
Multiflow	0.88	143	0.68	0.99	156	0.68	0.87	142	0.69
Multiflow w/o distillation	0.41	73	0.68	0.89	126	0.68	0.75	110	0.69
Multiflow w/o sequence	N/A	N/A	N/A	0.99	118	0.69	0.86	95	0.69

• DFMs は CTMC の確率性を介してサンプリング時の柔軟性を持つ離散データ生成を可能にし、適切な η に対してテキストデータで離散拡散ベースを上回る。
• Multiflow はタンパク質の構造と配列を共同生成するための最先端の共設計性能を達成。
• このモデルはどちらのモダリティにも条件付けをサポートし、インペインティングや順伝播/逆折りたたみタスクを含む柔軟な多模态サンプリングを提供。
• CTMC の確率性（η で制御）はサンプルの多様性とジャンプ頻度に影響を与え、探索と忠実性のバランスを取るノブを提供。
• 予備的な結果は Multiflow の逆折りたたみおよび前向き折りたたみの可能性を示唆しており、汎用的なタンパク質生成モデリングへの道を示している。