Сопоставление потоков Gumbel-Softmax с прямым сквозным управлением для контролируемой генерации биологических последовательностей
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
Аннотация
Сопоставление потоков в непрерывном симплексе стало перспективной стратегией для проектирования последовательностей ДНК, однако оно сталкивается с трудностями при масштабировании на более высокие размерности симплекса, необходимые для генерации пептидов и белков. Мы представляем Gumbel-Softmax Flow и Score Matching — генеративную структуру на симплексе, основанную на новом интерполянте Gumbel-Softmax с зависящей от времени температурой. Используя этот интерполянт, мы вводим Gumbel-Softmax Flow Matching, выводя параметризованное поле скоростей, которое переносит от гладких категориальных распределений к распределениям, сконцентрированным в одной вершине симплекса. Альтернативно мы представляем Gumbel-Softmax Score Matching, который обучается регрессии градиента плотности вероятности. Наша структура обеспечивает высококачественную и разнообразную генерацию и эффективно масштабируется на симплексы более высокой размерности. Для реализации обучения без обучения мы предлагаем Straight-Through Guided Flows (STGFlow) — метод управления на основе классификатора, который использует straight-through estimators для направления безусловного поля скоростей к оптимальным вершинам симплекса. STGFlow позволяет эффективно управлять на этапе вывода с использованием классификаторов, предварительно обученных на чистых последовательностях, и может применяться с любым методом дискретных потоков. Вместе эти компоненты формируют надежную структуру для управляемой генерации последовательностей de novo. Мы демонстрируем передовые результаты в условном проектировании промоторов ДНК, генерации белков только на основе последовательностей и проектировании пептидов для связывания с мишенями при лечении редких заболеваний.
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.