Repurposing Protein Language Models for Latent Flow-Based Fitness Optimization

TL;DR

Protein fitness optimization faces vast, sparse landscapes; CHASE delays no sampling by relying on a latent flow model conditioned on fitness. It repurposes pretrained protein language model embeddings into a compact latent space via a VAE and trains a conditional flow-matching model with classifier-free guidance to generate high-fitness sequences directly in latent space, reconstructing them back to sequences. The approach achieves state-of-the-art results on AAV and GFP benchmarks, and bootstrapping with synthetic data further improves performance in data-scarce regimes, while delivering substantial efficiency gains over predictor-guided methods. This framework demonstrates a scalable, predictor-free pathway for de novo protein design that preserves diversity and novelty while targeting high fitness.

Abstract

Protein fitness optimization is challenged by a vast combinatorial landscape where high-fitness variants are extremely sparse. Many current methods either underperform or require computationally expensive gradient-based sampling. We present CHASE, a framework that repurposes the evolutionary knowledge of pretrained protein language models by compressing their embeddings into a compact latent space. By training a conditional flow-matching model with classifier-free guidance, we enable the direct generation of high-fitness variants without predictor-based guidance during the ODE sampling steps. CHASE achieves state-of-the-art performance on AAV and GFP protein design benchmarks. Finally, we show that bootstrapping with synthetic data can further enhance performance in data-constrained settings.

Figures (4)

• Figure 1: Fitness–diversity–novelty chart across protein optimization methods for gfp Medium. method achieves high fitness while maintaining diversity and novelty compared to existing baselines, whereas alternative methods tend to favor individual objectives at the expense of others.
• Figure 2: Overview of method. (1) Discrete protein sequences are encoded by a pretrained and (2) projected into a compact latent manifold via a compression block. (3) A conditional generative model learns the probability path between noise and these embeddings, conditioned on time $t$ and fitness $f$. (4) At inference, high-fitness embeddings are sampled and mapped back to sequence space through a decompressor and a discrete decoder. Gray blocks denote modules trained in our framework.
• Figure 3: Detailed version of our architecture for the encoder/decoder setup, the compressor/decompressor and the flow matching model. The addition of the conditional score $f$ and the time conditioning $t$ was omitted for brevity.
• Figure 4: Effect of score dropout rate $p$ during training on model performance for AAV Hard and GFP Medium benchmarks.