Fluxo de Correspondência Gumbel-Softmax com Orientação Direta para Geração Controlável de Sequências 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

Resumo

O emparelhamento de fluxo no simplex contínuo surgiu como uma estratégia promissora para o design de sequências de DNA, mas enfrenta dificuldades para escalar para dimensões mais altas do simplex necessárias para a geração de peptídeos e proteínas. Introduzimos o Gumbel-Softmax Flow e Score Matching, um framework generativo no simplex baseado em um novo interpolante Gumbel-Softmax com uma temperatura dependente do tempo. Usando esse interpolante, apresentamos o Gumbel-Softmax Flow Matching ao derivar um campo de velocidade parametrizado que transporta de distribuições categóricas suaves para distribuições concentradas em um único vértice do simplex. Alternativamente, apresentamos o Gumbel-Softmax Score Matching, que aprende a regredir o gradiente da densidade de probabilidade. Nosso framework permite geração de alta qualidade e diversificada, escalando eficientemente para simplices de maior dimensão. Para permitir orientação sem treinamento, propomos o Straight-Through Guided Flows (STGFlow), um método de orientação baseado em classificadores que utiliza estimadores straight-through para direcionar o campo de velocidade incondicional em direção aos vértices ótimos do simplex. O STGFlow permite orientação eficiente no momento da inferência usando classificadores pré-treinados em sequências limpas e pode ser usado com qualquer método de fluxo discreto. Juntos, esses componentes formam um framework robusto para geração controlada de sequências de novo. Demonstramos desempenho de ponta no design condicional de promotores de DNA, geração de proteínas baseada apenas em sequências e design de peptídeos ligantes de alvos para o tratamento de doenças 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.