제어 가능한 생물학적 시퀀스 생성을 위한 스트레이트-스루 가이던스 기반 Gumbel-Softmax Flow Matching

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

March 21, 2025

저자: Sophia Tang, Yinuo Zhang, Alexander Tong, Pranam Chatterjee

cs.AI

초록

연속 심플렉스(continuous simplex)에서의 플로우 매칭(Flow Matching)은 DNA 서열 설계를 위한 유망한 전략으로 부상했지만, 펩타이드 및 단백질 생성에 필요한 더 높은 차원의 심플렉스로 확장하는 데 어려움을 겪고 있습니다. 우리는 시간에 따라 변하는 온도를 가진 새로운 Gumbel-Softmax 보간법을 기반으로 심플렉스 상에서의 생성 프레임워크인 Gumbel-Softmax Flow 및 스코어 매칭(Score Matching)을 소개합니다. 이 보간법을 사용하여, 우리는 매끄러운 범주형 분포에서 심플렉스의 단일 정점에 집중된 분포로 전달하는 매개변수화된 속도 필드를 도출함으로써 Gumbel-Softmax Flow Matching을 제안합니다. 또한, 우리는 확률 밀도의 기울기를 회귀하는 방법인 Gumbel-Softmax 스코어 매칭을 제안합니다. 우리의 프레임워크는 고품질의 다양한 생성을 가능하게 하며, 더 높은 차원의 심플렉스로 효율적으로 확장할 수 있습니다. 훈련 없이도 가이던스를 가능하게 하기 위해, 우리는 무조건적 속도 필드를 심플렉스의 최적 정점으로 조종하기 위해 스트레이트-스루 추정기(straight-through estimators)를 활용하는 분류기 기반 가이던스 방법인 Straight-Through Guided Flows(STGFlow)를 제안합니다. STGFlow는 깨끗한 서열에 대해 사전 훈련된 분류기를 사용하여 추론 시간에 효율적인 가이던스를 가능하게 하며, 모든 이산 플로우 방법과 함께 사용할 수 있습니다. 이러한 구성 요소들은 통제 가능한 데노보(de novo) 서열 생성을 위한 강력한 프레임워크를 형성합니다. 우리는 조건부 DNA 프로모터 설계, 서열만을 이용한 단백질 생성, 희귀 질환 치료를 위한 표적 결합 펩타이드 설계에서 최첨단 성능을 입증합니다.

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.