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

Flow matching in discrete spaces is powerful but limited by high variance (Dirichlet FM), potential overfitting (Fisher FM), and rigidity. We build on this by using a Gumbel-Softmax interpolant with time-dependent temperature to create smooth, learnable flows across the simplex! 🔥

Our model, Gumbel-Softmax FM, smoothly transforms noise into clean sequences—avoiding hard discretization while learning better transport paths. We also introduce Gumbel-Softmax Score Matching, learning the score function over the simplex for stochastic sampling! 🎯

But generative control is hard post-training. 😓 So we introduce STGFlow — a training-free guidance method using straight-through gradients from pre-trained classifiers to steer flow trajectories at inference. No need to retrain: plug-and-play guidance for peptides, DNA, and proteins!!🧠

Across bioengineering tasks—DNA promoter design, de novo protein generation, and peptide binder design—our framework outperforms autoregressive and discrete diffusion baselines in fidelity, foldability, and functional control.📈🧬

We even show de novo peptide binders to targets with no known binders, including proteins involved in rare pediatric leukodystrophies and neurodegenerative diseases.💊 Our binders further show better docking scores and ipTM values than known binders across 13 targets and scramble controls! 🙌