Flux d'appariement Gumbel-Softmax avec guidage direct pour la génération contrôlée de séquences biologiques

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

Résumé

L'appariement de flux dans le simplexe continu s'est imposé comme une stratégie prometteuse pour la conception de séquences d'ADN, mais peine à s'adapter aux dimensions plus élevées du simplexe nécessaires à la génération de peptides et de protéines. Nous introduisons Gumbel-Softmax Flow et Score Matching, un cadre génératif sur le simplexe basé sur un nouvel interpolant Gumbel-Softmax avec une température dépendante du temps. En utilisant cet interpolant, nous introduisons Gumbel-Softmax Flow Matching en dérivant un champ de vitesse paramétré qui transporte des distributions catégorielles lisses vers des distributions concentrées sur un seul sommet du simplexe. Nous présentons également Gumbel-Softmax Score Matching, qui apprend à régresser le gradient de la densité de probabilité. Notre cadre permet une génération de haute qualité et diversifiée, et s'adapte efficacement aux simplexes de plus haute dimension. Pour permettre un guidage sans entraînement, nous proposons Straight-Through Guided Flows (STGFlow), une méthode de guidage basée sur des classifieurs qui exploite des estimateurs straight-through pour orienter le champ de vitesse inconditionnel vers les sommets optimaux du simplexe. STGFlow permet un guidage efficace à l'inférence en utilisant des classifieurs pré-entraînés sur des séquences propres, et peut être utilisé avec n'importe quelle méthode de flux discret. Ensemble, ces composants forment un cadre robuste pour la génération contrôlée de séquences de novo. Nous démontrons des performances de pointe dans la conception conditionnelle de promoteurs d'ADN, la génération de protéines basée uniquement sur les séquences, et la conception de peptides liant des cibles pour le traitement de maladies rares.

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.