Flusso di Gumbel-Softmax con Guida Diretta per la Generazione Controllabile di Sequenze Biologiche

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

March 21, 2025

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

cs.AI

Abstract

Il matching dei flussi nel simplesso continuo è emerso come una strategia promettente per la progettazione di sequenze di DNA, ma fatica a scalare verso dimensioni più elevate del simplesso necessarie per la generazione di peptidi e proteine. Introduciamo il Gumbel-Softmax Flow e il Score Matching, un framework generativo sul simplesso basato su un nuovo interpolante Gumbel-Softmax con una temperatura dipendente dal tempo. Utilizzando questo interpolante, introduciamo il Gumbel-Softmax Flow Matching derivando un campo di velocità parametrizzato che trasporta da distribuzioni categoriche lisce a distribuzioni concentrate su un singolo vertice del simplesso. Alternativamente, presentiamo il Gumbel-Softmax Score Matching che apprende a regredire il gradiente della densità di probabilità. Il nostro framework consente una generazione di alta qualità e diversificata, scalando efficientemente verso simplessi di dimensioni superiori. Per abilitare una guida senza addestramento, proponiamo gli Straight-Through Guided Flows (STGFlow), un metodo di guida basato su classificatori che sfrutta stimatori straight-through per indirizzare il campo di velocità incondizionato verso i vertici ottimali del simplesso. STGFlow consente una guida efficiente al momento dell'inferenza utilizzando classificatori pre-addestrati su sequenze pulite, e può essere utilizzato con qualsiasi metodo di flusso discreto. Insieme, questi componenti formano un framework robusto per la generazione controllata di sequenze de novo. Dimostriamo prestazioni all'avanguardia nella progettazione condizionale di promotori di DNA, nella generazione di proteine basata esclusivamente su sequenze, e nella progettazione di peptidi leganti per il trattamento di malattie rare.

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.