Flows, straight but not so fast: Exploring the design space of Rectified Flows in Protein Design

TL;DR

This paper explores how Rectified Flows can be adapted for protein backbone generation, reducing computational costs and improving efficiency compared to traditional diffusion models.

Contribution

It systematically studies ReFlow design choices for proteins, demonstrating their sensitivity and proposing improvements specific to protein generation tasks.

Findings

ReFlow reduces the number of function evaluations needed for protein generation.

Design choices from image ReFlow do not directly apply to proteins.

Proposed methodological improvements enhance ReFlow performance for proteins.

Abstract

Generative modeling techniques such as Diffusion and Flow Matching have achieved significant successes in generating designable and diverse protein backbones. However, many current models are computationally expensive, requiring hundreds or even thousands of function evaluations (NFEs) to yield samples of acceptable quality, which can become a bottleneck in practical design campaigns that often generate $10^4\ -\ 10^6$ designs per target. In image generation, Rectified Flows (ReFlow) can significantly reduce the required NFEs for a given target quality, but their application in protein backbone generation has been less studied. We apply ReFlow to improve the low NFE performance of pretrained SE(3) flow matching models for protein backbone generation and systematically study ReFlow design choices in the context of protein generation in data curation, training and inference time settings. In particular, we (1) show that ReFlow in the protein domain is particularly sensitive to the choice of coupling generation and annealing, (2) demonstrate how useful design choices for ReFlow in the image domain do not directly translate to better performance on proteins, and (3) make improvements to ReFlow methodology for proteins.