Gumbel-Softmax Flow Matching mit Straight-Through Guidance für die kontrollierbare Erzeugung biologischer Sequenzen

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

March 21, 2025

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

cs.AI

Zusammenfassung

Flow Matching im kontinuierlichen Simplex hat sich als vielversprechende Strategie für das Design von DNA-Sequenzen erwiesen, stößt jedoch bei der Skalierung auf höhere Simplex-Dimensionen, die für die Peptid- und Proteingenerierung erforderlich sind, an Grenzen. Wir stellen Gumbel-Softmax Flow und Score Matching vor, ein generatives Framework auf dem Simplex, das auf einem neuartigen Gumbel-Softmax-Interpolanten mit einer zeitabhängigen Temperatur basiert. Mit diesem Interpolanten führen wir Gumbel-Softmax Flow Matching ein, indem wir ein parametrisiertes Geschwindigkeitsfeld ableiten, das von glatten kategorialen Verteilungen zu Verteilungen transportiert, die an einem einzelnen Scheitelpunkt des Simplex konzentriert sind. Alternativ präsentieren wir Gumbel-Softmax Score Matching, das lernt, den Gradienten der Wahrscheinlichkeitsdichte zu regressieren. Unser Framework ermöglicht eine hochwertige, vielfältige Generierung und skaliert effizient auf höherdimensionale Simplizes. Um ein trainingsfreies Guidance zu ermöglichen, schlagen wir Straight-Through Guided Flows (STGFlow) vor, eine klassifikatorbasierte Guidance-Methode, die Straight-Through-Schätzer nutzt, um das unbedingte Geschwindigkeitsfeld in Richtung optimaler Scheitelpunkte des Simplex zu lenken. STGFlow ermöglicht eine effiziente Inferenzzeit-Guidance mit auf sauberen Sequenzen vortrainierten Klassifikatoren und kann mit jeder diskreten Flow-Methode verwendet werden. Zusammen bilden diese Komponenten ein robustes Framework für die kontrollierbare de novo Sequenzgenerierung. Wir demonstrieren state-of-the-art Leistungen im konditionalen Design von DNA-Promotoren, der sequenzbasierten Proteingenerierung und dem Design von zielbindenden Peptiden für die Behandlung seltener Krankheiten.

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.