Gumbel-Softmax Flow Matching met Straight-Through Guidance voor Beheersbare Generatie van Biologische Sequenties

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

March 21, 2025

Auteurs: Sophia Tang, Yinuo Zhang, Alexander Tong, Pranam Chatterjee

cs.AI

Samenvatting

Flow matching in het continue simplex is naar voren gekomen als een veelbelovende strategie voor DNA-sequentieontwerp, maar heeft moeite om op te schalen naar hogere simplexdimensies die nodig zijn voor peptide- en proteïnegeneratie. We introduceren Gumbel-Softmax Flow en Score Matching, een generatief framework op het simplex gebaseerd op een nieuwe Gumbel-Softmax-interpolant met een tijdsafhankelijke temperatuur. Met behulp van deze interpolant introduceren we Gumbel-Softmax Flow Matching door een geparametriseerd snelheidsveld af te leiden dat transporteert van gladde categorische verdelingen naar verdelingen geconcentreerd op een enkel hoekpunt van het simplex. We presenteren alternatief Gumbel-Softmax Score Matching, dat leert om de gradiënt van de waarschijnlijkheidsdichtheid te regresseren. Ons framework maakt hoogwaardige, diverse generatie mogelijk en schaalt efficiënt naar hoger-dimensionale simplices. Om training-vrije begeleiding mogelijk te maken, stellen we Straight-Through Guided Flows (STGFlow) voor, een classifier-gebaseerde begeleidingsmethode die straight-through estimators benut om het onvoorwaardelijke snelheidsveld te sturen naar optimale hoekpunten van het simplex. STGFlow maakt efficiënte inferentie-tijdbegeleiding mogelijk met classifiers die vooraf zijn getraind op schone sequenties, en kan worden gebruikt met elke discrete flow-methode. Samen vormen deze componenten een robuust framework voor controleerbare de novo-sequentiegeneratie. We demonstreren state-of-the-art prestaties in conditioneel DNA-promoterontwerp, sequentie-alleen proteïnegeneratie en doelbindend peptideontwerp voor de behandeling van zeldzame ziekten.

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.