Gumbel-Softmax Flow MatchingとStraight-Throughガイダンスを用いた制御可能な生物学的配列生成

Gumbel-Softmax Flow Matching with Straight-Through Guidance for Controllable Biological Sequence Generation

March 21, 2025

著者: Sophia Tang, Yinuo Zhang, Alexander Tong, Pranam Chatterjee

cs.AI

要旨

連続シンプレックスにおけるフローマッチングは、DNA配列設計の有望な戦略として登場しましたが、ペプチドやタンパク質生成に必要な高次元シンプレックスへのスケーリングが困難です。本論文では、時間依存の温度を持つ新しいGumbel-Softmax補間法に基づくシンプレックス上の生成フレームワークであるGumbel-Softmax Flow and Score Matchingを提案します。この補間法を用いて、滑らかなカテゴリ分布からシンプレックスの単一頂点に集中した分布へと輸送するパラメータ化された速度場を導出し、Gumbel-Softmax Flow Matchingを導入します。また、確率密度の勾配を回帰するGumbel-Softmax Score Matchingも提示します。本フレームワークは、高品質で多様な生成を可能とし、高次元シンプレックスへの効率的なスケーリングを実現します。訓練不要のガイダンスを可能にするため、Straight-Through Guided Flows (STGFlow)を提案します。これは、ストレートスルー推定量を活用して無条件の速度場をシンプレックスの最適頂点へと誘導する分類器ベースのガイダンス手法です。STGFlowは、クリーンな配列で事前訓練された分類器を用いた効率的な推論時ガイダンスを可能とし、任意の離散フロー手法と組み合わせて使用できます。これらのコンポーネントを組み合わせることで、制御可能なデノボ配列生成のための堅牢なフレームワークを構築します。本手法は、条件付きDNAプロモーター設計、配列のみに基づくタンパク質生成、希少疾患治療のためのターゲット結合ペプチド設計において、最先端の性能を実証します。

English

Flow matching in the continuous simplex has emerged as a promising strategy for DNA sequence design, but struggles to scale to higher simplex dimensions required for peptide and protein generation. We introduce Gumbel-Softmax Flow and Score Matching, a generative framework on the simplex based on a novel Gumbel-Softmax interpolant with a time-dependent temperature. Using this interpolant, we introduce Gumbel-Softmax Flow Matching by deriving a parameterized velocity field that transports from smooth categorical distributions to distributions concentrated at a single vertex of the simplex. We alternatively present Gumbel-Softmax Score Matching which learns to regress the gradient of the probability density. Our framework enables high-quality, diverse generation and scales efficiently to higher-dimensional simplices. To enable training-free guidance, we propose Straight-Through Guided Flows (STGFlow), a classifier-based guidance method that leverages straight-through estimators to steer the unconditional velocity field toward optimal vertices of the simplex. STGFlow enables efficient inference-time guidance using classifiers pre-trained on clean sequences, and can be used with any discrete flow method. Together, these components form a robust framework for controllable de novo sequence generation. We demonstrate state-of-the-art performance in conditional DNA promoter design, sequence-only protein generation, and target-binding peptide design for rare disease treatment.