Repurposing Protein Language Models for Latent Flow-Based Fitness Optimization

TL;DR

CHASE framework leverages pretrained protein language models and flow matching to generate high-fitness protein variants efficiently without external predictors.

cs.LG 🔴 Advanced 2026-02-03 40 views
Amaru Caceres Arroyo Lea Bogensperger Ahmed Allam Michael Krauthammer Konrad Schindler Dominik Narnhofer
protein engineering deep learning generative models latent space optimization

Key Findings

Methodology

This study employs pretrained protein language models (like ESM2) to extract high-dimensional embeddings of protein sequences, which are then compressed into a low-dimensional latent space using a variational autoencoder (VAE). A conditional flow matching (Flow Matching) model is trained to learn the distribution of these latent representations conditioned on fitness scores, utilizing classifier-free guidance to enable efficient, target-directed sampling. During inference, the model generates high-fitness sequences by sampling from the learned conditional distribution in latent space, which are then decoded back into amino acid sequences. Synthetic data augmentation is employed to address data scarcity, enhancing the model's ability to explore the fitness landscape effectively.

Key Results

  • On AAV and GFP benchmarks, CHASE outperforms existing methods, achieving median fitness scores of 0.53 (AAV Medium), surpassing GGS (0.54) and VLGPO (0.48). It maintains competitive diversity and novelty metrics, demonstrating robust exploration. The approach reduces inference time by factors of 10-85 compared to gradient-based methods like VLGPO, thanks to classifier-free guidance, which eliminates the need for costly predictor evaluations during sampling.
  • Synthetic data bootstrapping further improves performance in low-data regimes, enabling the model to generate sequences with higher fitness and diversity. The experiments validate that the latent space conditioning approach effectively balances exploration and exploitation, leading to state-of-the-art results across multiple benchmarks.
  • Overall, CHASE demonstrates that leveraging pretrained embeddings and flow-based generative modeling can significantly advance protein design, offering a scalable, efficient alternative to traditional gradient-guided methods.

Significance

This work addresses the core challenge of exploring vast, discrete protein sequence spaces efficiently. By transforming the problem into a continuous latent space guided by pretrained knowledge, it overcomes the limitations of gradient-based and predictor-dependent methods, reducing computational costs while enhancing diversity and target fitness. The approach paves the way for scalable, automated protein engineering applicable in drug discovery, enzyme optimization, and synthetic biology, potentially accelerating development cycles and reducing experimental costs. Its ability to generate diverse, high-fitness variants without external predictors marks a significant step forward in AI-driven molecular design.

Technical Contribution

The paper introduces an innovative integration of pretrained protein language model embeddings, variational autoencoders for latent compression, and flow matching with classifier-free guidance for conditional sequence generation. This combination enables fast, target-specific protein design without reliance on external fitness predictors. The model's architecture allows for efficient sampling, synthetic data augmentation, and improved exploration of the fitness landscape. Theoretical guarantees include stable training of the flow model and effective conditioning in latent space, leading to superior performance over prior methods that depend on predictor-guided sampling or operate directly in discrete sequence space.

Novelty

This is the first work to embed pretrained protein language model representations into a compressed latent space for conditional flow-based generation, avoiding the need for external predictors during sampling. Unlike prior approaches that operate directly on sequences or rely on gradient guidance, CHASE leverages the rich prior knowledge encoded in large-scale pLMs, combined with flow models, to enable efficient, target-directed protein design. Its synthetic data augmentation strategy further enhances exploration in data-scarce scenarios, setting a new standard for scalable, predictor-free protein optimization.

Limitations

  • The approach's effectiveness depends heavily on the quality of the pretrained embeddings; biases or gaps in the pLM training data can limit performance. In extremely sparse or biased datasets, the model may struggle to generalize.
  • Potential information loss during latent compression could reduce the diversity or fine-grained control over generated sequences.
  • Computational costs remain significant, especially during training of the flow model and data augmentation phases, which may limit scalability in some settings.

Future Work

Future directions include integrating structural and biophysical information into the latent space to improve functional relevance, exploring reinforcement learning for more adaptive optimization, and extending the framework to multi-objective tasks such as stability and solubility. Additionally, combining experimental feedback with synthetic data generation could further refine the model's accuracy and applicability in real-world protein engineering workflows.

AI Executive Summary

Designing proteins with desired functions remains a central challenge in molecular biology. Traditional methods like directed evolution, while effective, are resource-intensive and limited in exploring the enormous sequence space. Recent advances in deep learning, especially pretrained protein language models such as ESM2, have revolutionized the field by capturing evolutionary and structural priors from massive sequence datasets. These models generate rich embeddings that encode complex relationships, offering a promising foundation for computational protein design.

However, leveraging these embeddings for targeted optimization requires overcoming the discrete and high-dimensional nature of protein sequences. Existing approaches often depend on gradient-based sampling or external predictors, which are computationally costly and prone to overfitting. To address this, the authors propose CHASE, a novel framework that compresses pretrained embeddings into a continuous latent space using variational autoencoders. Within this space, a flow matching model is trained to learn the distribution of sequences conditioned on desired fitness levels, employing classifier-free guidance to enable efficient, predictor-free sampling.

This approach allows for direct, fast generation of high-fitness protein variants by integrating the learned conditional distribution. The method was tested on benchmark datasets for AAV and GFP proteins, demonstrating superior performance in fitness scores while maintaining diversity and novelty. Notably, the inference process is significantly faster—by factors of 10 to 85—compared to gradient-guided methods like VLGPO, making it highly scalable.

Furthermore, the authors introduce synthetic data augmentation strategies, which improve performance in data-limited scenarios, highlighting the robustness of the framework. Overall, CHASE represents a major step forward in AI-driven protein engineering, combining the power of pretrained models, latent space modeling, and flow-based generative techniques to enable efficient, scalable, and high-quality protein design. Its implications extend to drug discovery, enzyme engineering, and synthetic biology, promising to accelerate innovation in these fields.

Deep Analysis

Background

蛋白质工程经历了从传统的实验筛选到深度学习辅助设计的转变。早期方法如定向进化依赖大量实验,成本高昂且效率有限。近年来,预训练的蛋白质语言模型(如ESM、ProtTrans)通过学习大规模序列库,捕获了蛋白质的结构和功能信息,为自动化设计提供了基础。尽管如此,离散序列空间的高维稀疏性仍限制了优化效率。研究者尝试引入连续潜在空间(如VAE、流模型)以平滑探索,但多依赖外部预测器或梯度指导,存在计算成本高、探索受限等问题。本文在此背景下,结合预训练模型的丰富先验,提出一种高效的目标导向生成方法,旨在突破现有瓶颈。

Core Problem

蛋白质序列空间庞大且离散,导致优化难度极高。传统方法如随机突变和筛选耗费巨大,难以在有限样本下找到高性能变体。现有深度生成模型多依赖外部预测器或梯度引导,计算成本高且易陷入局部最优。此外,数据稀疏和偏差严重限制了模型的泛化能力。如何在保持多样性的同时,快速、准确地生成目标性能的蛋白质,成为亟待解决的核心问题。

Innovation

本研究的创新点包括:1)利用预训练蛋白质语言模型的嵌入,压缩到潜在空间,增强表达能力;2)引入条件流匹配模型,结合无分类引导,实现高效目标导向采样;3)避免外部预测器,降低计算成本,提升采样速度;4)采用合成数据增强策略,缓解数据稀疏问题。这些创新结合,显著提升了蛋白质设计的效率和效果,突破了离散空间探索的瓶颈,为自动化蛋白质工程提供了新思路。

Methodology

  • �� 利用预训练的ESM2模型提取蛋白序列的高维嵌入。• 通过变分自编码器(VAE)将嵌入压缩到低维潜在空间。• 训练条件流匹配模型,学习潜在空间中序列的分布,条件为目标适应性值。• 在采样阶段,利用无分类引导的流模型,从噪声开始,逐步生成高适应性潜在向量。• 将潜在向量解码回序列空间,得到候选蛋白。• 采用合成数据增强,通过在目标适应性区间生成虚拟样本,丰富训练集。• 最终,结合合成数据,反复训练模型,提升性能。

Experiments

采用AAV和GFP的基准数据集,训练样本有限(如AAV Medium仅有2139个样本)。模型性能通过适应性得分、多样性和新颖性指标评估。与GFlowNet、CbAS、AdaLead等多种方法对比,验证CHASE在适应性和探索能力上的优势。采用不同随机种子多次实验,确保结果稳定。关键超参数包括潜在空间维度、流模型的层数和采样步数。还测试了合成数据增强策略的效果,观察其对低样本环境的提升作用。

Results

CHASE在所有基准测试中均表现优异,适应性得分明显高于对比方法,例如在AAV Medium中达0.53,优于GGS的0.54和VLGPO的0.48。多样性指标保持在合理范围内,且新颖性略优于其他方法。无分类引导采样显著降低了计算成本,采样速度提升10-85倍。合成数据增强进一步改善了低样本条件下的性能,验证了策略的有效性。整体结果表明,潜在空间条件生成在蛋白质设计中具有巨大潜力。

Applications

该方法可应用于药物蛋白设计、酶工程、抗体优化等领域,特别适合有限样本或高成本实验场景。通过快速生成高性能候选,缩短研发周期,降低成本。未来结合结构信息和实验验证,可实现更复杂的功能设计,推动工业化应用。

Limitations & Outlook

模型在极端高维或极端稀疏的适应性空间中可能表现不佳,尤其是在训练数据极少或偏差较大的情况下。潜在空间的压缩可能导致信息损失,影响生成序列的多样性和细粒度调控能力。目前的架构仍依赖预训练模型的质量,若预训练模型存在偏差或不足,可能影响最终性能。

Plain Language Accessible to non-experts

想象你在厨房里做菜,手里有很多不同的食材(蛋白质序列),每种食材都有不同的味道和效果(功能)。传统做菜需要试错,花费很多时间和材料才能找到好吃的菜(高性能蛋白)。现在,有了一个智能厨师(模型),它已经学习了很多菜谱(预训练模型),知道哪些食材搭配会更好。这个厨师可以把复杂的食材组合(高维嵌入)压缩成简单的配方(潜在空间),然后根据你想要的味道(目标适应性),快速生成新菜谱(蛋白序列)。它不用试错,也不用反复试验,只需在脑海中模拟几步,就能做出美味的菜。这就像用智能算法提前规划好做菜步骤,节省时间,又能做出多样化的菜肴,满足不同口味需求。

ELI14 Explained like you're 14

想象你在学校的烹饪课上,老师给你一些食材(蛋白质序列),让你做出既好吃又特别的菜(高适应性蛋白)。以前,你只能靠自己试错,试了很多次才能找到好吃的组合,既费时间又浪费材料。现在,有了一个超级聪明的机器人厨师(模型),它已经学会了很多菜谱(预训练模型),可以帮你设计出新菜。这个机器人把复杂的食材搭配(高维信息)变成简单的配方(潜在空间),你只要告诉它你想要的味道(目标适应性),它就能快速帮你做出多样的菜。这个过程就像用智能程序提前规划好做菜步骤,不用反复试验,就能做出美味又新颖的菜肴。这样,你就可以更快、更轻松地做出各种美味的菜,满足不同的口味需求。

Glossary

Protein Language Model(蛋白质语言模型)

一种深度学习模型,学习大量蛋白序列的结构和功能信息,能生成或理解蛋白质序列。In this paper, used to提取蛋白序列的嵌入特征。

作为特征提取和先验知识的基础模型。

Flow Matching(流匹配)

一种生成模型,通过学习时间依赖的向量场,将噪声逐步变换为目标分布。In this研究中,用于潜在空间的条件采样。

实现高效目标导向的蛋白质序列生成。

Variational Autoencoder(变分自编码器)

一种深度生成模型,将高维数据映射到低维潜在空间,兼顾重建和正则化。用于蛋白质嵌入的压缩。

在潜在空间中进行连续优化。

Classifier-Free Guidance(无分类引导)

一种条件生成技术,通过在训练中随机丢弃条件信息,提升模型的采样效率和多样性。

在采样阶段引导目标导向生成。

Synthetic Data Augmentation(合成数据增强)

利用模型生成虚拟样本,扩展训练集,缓解数据稀疏问题。

提升模型在有限样本下的性能。

Open Questions Unanswered questions from this research

  • 1 如何进一步结合蛋白质结构信息以提升生成的功能性和稳定性仍未解决,尤其是在结构-功能关系复杂的蛋白中。
  • 2 潜在空间的压缩可能导致信息损失,如何平衡表达能力与信息保留是未来研究重点。

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.

cs.LG q-bio.QM