[论文解读] ReLSO: A Transformer-based Model for Latent Space Optimization and Generation of Proteins

Egbert Castro, Abhinav Godavarthi+4 more|arXiv (Cornell University)|Jan 24, 2022

vaccines and immunoinformatics approaches被引用 6

一句话总结

ReLSO 是一种基于 Transformer 的自编码器，可学习蛋白质序列的结构化、正则化潜在空间，用于联合进行蛋白质序列生成与功能预测。通过在潜在空间中实现基于梯度的优化，ReLSO 在抗-ranibizumab 和 GFP 数据集上实现了比基线方法更高的功能优化效率，并生成了更具鲁棒性的高功能序列。

ABSTRACT

The development of powerful natural language models have increased the ability to learn meaningful representations of protein sequences. In addition, advances in high-throughput mutagenesis, directed evolution, and next-generation sequencing have allowed for the accumulation of large amounts of labeled fitness data. Leveraging these two trends, we introduce Regularized Latent Space Optimization (ReLSO), a deep transformer-based autoencoder which features a highly structured latent space that is trained to jointly generate sequences as well as predict fitness. Through regularized prediction heads, ReLSO introduces a powerful protein sequence encoder and novel approach for efficient fitness landscape traversal. Using ReLSO, we explicitly model the sequence-function landscape of large labeled datasets and generate new molecules by optimizing within the latent space using gradient-based methods. We evaluate this approach on several publicly-available protein datasets, including variant sets of anti-ranibizumab and GFP. We observe a greater sequence optimization efficiency (increase in fitness per optimization step) by ReLSO compared to other approaches, where ReLSO more robustly generates high-fitness sequences. Furthermore, the attention-based relationships learned by the jointly-trained ReLSO models provides a potential avenue towards sequence-level fitness attribution information.

研究动机与目标

• 解决在庞大且存在上位效应的蛋白质序列空间中高效设计功能性蛋白质的挑战。
• 利用高通量筛选获得的大规模带标签功能数据，训练一个联合建模序列与功能的深度生成模型。
• 构建一个高度结构化、平滑且伪凸的潜在空间，以实现蛋白质序列的高效基于梯度的优化。
• 通过在联合训练的功能预测头上的梯度上升，实现在潜在空间中的直接优化，从而在黑箱或顺序方法之上提升搜索效率。

提出的方法

• 采用基于 Transformer 的自编码器架构，将蛋白质序列编码为低维、信息丰富的潜在空间。
• 通过正则化功能预测头联合训练潜在表征，直接从潜在码预测功能特性。
• 应用三种关键正则化策略：(1) 通过功能预测损失实现平滑性，(2) 通过潜在空间正则化实现连续性与可插值性，(3) 通过训练数据外的负采样实现伪凸性。
• 利用功能预测头的梯度，在潜在空间中执行基于梯度的优化，实现高效向高功能区域的遍历。
• 将 ReLSO 的性能与多种基线方法进行比较，包括体外定向进化、MCMC、爬山法以及自适应采样（DbAS、CbAS），涵盖序列空间与潜在空间。
• 使用公开的 DbAS 与 CbAS 实现，并通过超参数调优（q ∈ [0.5, 0.8]，epochs ∈ [1, 15]）确保公平比较。

实验结果

研究问题

• RQ1 基于 Transformer 的自编码器中，联合训练并正则化的潜在空间是否能实现比标准黑箱或顺序方法更高效的蛋白质序列优化？
• RQ2 ReLSO 的正则化潜在空间在多大程度上提升了功能景观的遍历性，并加速收敛至高功能序列？
• RQ3 在每一步优化中，潜在空间中的基于梯度优化与 MCMC、爬山法及自适应采样等无梯度方法相比，其功能增益如何？
• RQ4 ReLSO 中的注意力机制是否通过学习到的注意力模式，提供了可解释的、序列级别的功能归因信息？
• RQ5 ReLSO 是否能在不同蛋白质家族（如抗-ranibizumab 和 GFP）之间泛化，实现仅需极少体外筛选即可生成高功能变体？

主要发现

• 在抗-ranibizumab 和 GFP 数据集上，ReLSO 显著提升了功能优化效率——以每步优化的平均功能提升衡量——优于其他方法。
• ReLSO 模型能稳健地生成高功能序列，其中基于梯度的潜在空间优化（ReLSO-GA）优于无梯度方法及黑箱自适应采样方法。
• 联合训练过程中学习到的基于注意力的关系，为序列级别的功能归因提供了潜在路径，实现了超越简单预测的可解释性。
• 正则化的潜在空间支持平滑、连续且可插值的表征，有助于稳定且高效的梯度上升，降低陷入局部最优的风险。
• 由于潜在空间中集成了结构与功能归纳偏置，ReLSO 的性能优于最先进的方法（如 DbAS 与 CbAS），尤其在早期优化阶段表现更优。
• 负采样与插值正则化增强了潜在空间的伪凸性，提升了优化稳定性与收敛性。