Gumbel-Softmax Flow Matching with Straight-Through Guidance for Controllable Biological Sequence Generation (2503.17361v1)

Abstract: 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.

Gumbel-Softmax Flow Matching for Biological Sequence Generation

The paper presents Gumbel-Softmax Flow Matching (Gumbel-Softmax FM) and Score Matching (Gumbel-Softmax SM), a framework designed to generate discrete biological sequences efficiently via a novel interpolant grounded on the Gumbel-Softmax distribution. The research targets the challenge of scalable sequence generation, focusing on achieving diverse and high-quality outputs, particularly in generating DNA, protein, and peptide sequences. It proposes a robust methodology to navigate the simplex, maintaining computational efficiency and flexibility independent of complex training protocols typically necessary for effective de novo design in sequencing tasks.

Methodology and Contributions

The core innovation lies in the introduction of a Gumbel-Softmax interpolant with a time-dependent temperature parameter that orchestrates the transformation of noisy categorical distributions toward clean, concentrated outputs. The paper establishes two core models in its framework:

1. Gumbel-Softmax Flow Matching (FM): It leverages the derived interpolant to define a velocity field that facilitates smooth distribution transport across the simplex, crucial for generating sequences that align well with desired structural properties.
2. Gumbel-Softmax Score Matching (SM): This component estimates the gradient of the probability density explicit to high-density simplex regions, thereby refining the sampling process from these regions.

A key technical contribution is the application of Straight-Through Guidance using STGFlow. This mechanism allows efficient classifier-based inference-time adjustments to streamline the generated sequence distribution even without extensive post-training, ensuring controllability and optimizing for desired binding affinities in peptide sequences.

Results and Evaluation

The results substantiate the framework’s capability to outperform existing methods both in efficiency and quality of outputs. Comparative experiments demonstrate superior performance in DNA promoter design, where the proposed model excels in signal prediction metrics. For de novo protein sequence generation, the research showcases competitive structural quality compared to leading contemporary models like EvoDiff and ProtGPT2, while maintaining relative simplicity in the model parameterization. Furthermore, in peptide binder design, the integration of classifier-based guidance markedly enhances output affinity, with new peptide designs showcasing improved binding scores over known sequences.

Implications and Future Directions

This paper introduces a promising pathway toward scalable and controllable sequence generation focused on biological applications. The strategic use of the Gumbel-Softmax distribution mitigates discretization errors, frequently encountered in categorical generation processes in higher-dimensional spaces. The integration of STGFlow illustrates a step forward in post-training optimization strategies that could find applications in broader AI models dealing with discrete data management.

Future work may explore the extension of this framework to facilitate multi-objective sequence optimizations, such as jointly optimizing for structural stability and functional activity. As the landscape of AI-driven biological design advances, particularly in drug discovery and personalized medicine, methodologies that couple efficient, robust generation and seamless post-generation controls like those proposed in this paper will be central to practical implementations and innovations.

By eliminating stringent training dependencies and integrating a flexible guidance strategy, this framework sets a new precedent for simplex-based generative models, offering a versatile tool for continual developments in biomolecular and genomic research.