Flujo de Gumbel-Softmax con Guía de Paso Directo para la Generación Controlada de Secuencias Biológicas

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

March 21, 2025

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

cs.AI

Resumen

El emparejamiento de flujos en el símplex continuo ha surgido como una estrategia prometedora para el diseño de secuencias de ADN, pero enfrenta dificultades para escalar a dimensiones más altas del símplex requeridas para la generación de péptidos y proteínas. Introducimos Gumbel-Softmax Flow y Score Matching, un marco generativo en el símplex basado en un nuevo interpolante Gumbel-Softmax con una temperatura dependiente del tiempo. Utilizando este interpolante, presentamos Gumbel-Softmax Flow Matching al derivar un campo de velocidad parametrizado que transporta desde distribuciones categóricas suaves hacia distribuciones concentradas en un solo vértice del símplex. Alternativamente, presentamos Gumbel-Softmax Score Matching, que aprende a regresar el gradiente de la densidad de probabilidad. Nuestro marco permite una generación de alta calidad y diversa, y escala eficientemente a símplices de mayor dimensión. Para habilitar la guía sin entrenamiento, proponemos Straight-Through Guided Flows (STGFlow), un método de guía basado en clasificadores que aprovecha estimadores straight-through para dirigir el campo de velocidad incondicional hacia los vértices óptimos del símplex. STGFlow permite una guía eficiente en tiempo de inferencia utilizando clasificadores preentrenados en secuencias limpias, y puede usarse con cualquier método de flujo discreto. Juntos, estos componentes forman un marco robusto para la generación controlada de secuencias de novo. Demostramos un rendimiento de vanguardia en el diseño condicional de promotores de ADN, la generación de proteínas basada únicamente en secuencias, y el diseño de péptidos de unión a objetivos para el tratamiento de enfermedades raras.

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.