Gumbel-Softmax Flow Matching with Straight-Through Guidance for
  Controllable Biological Sequence Generation
Authors: name: Sophia Tang; name: Yinuo Zhang; name: Alexander Tong; name: Pranam Chatterjee
Published: 2025-03-21 17:59:43

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.