One-sentence Summary

Researchers from Shanghai Jiao Tong University, East China University of Science and Technology, and Shanghai Artificial Intelligence Laboratory propose VENUSREM, a retrieval-enhanced protein language model that captures local amino acid interactions in spatial and temporal scales, achieving state-of-the-art performance on 217 assays from the ProteinGym benchmark and validated through wet-lab experiments demonstrating improved stability and binding affinity of a VHH antibody via post hoc analysis of over 30 mutants and enhanced activity of 10 designed DNA polymerase mutants at elevated temperatures, thereby establishing a comprehensive evaluation framework for mutation effect prediction.

Key Contributions

• The paper presents VENUSREM, a retrieval-enhanced protein language model that integrates evolutionary information from homologous sequences to capture local amino acid interactions across spatial and temporal scales, achieving state-of-the-art performance on 217 ProteinGym assays.
• This work establishes a multidimensional post hoc analysis scheme using low-throughput experimental data to complement high-throughput benchmarks, enabling independent model validation without new wet-lab experiments; it verifies improved stability and binding affinity for over 30 VHH antibody mutants.
• The method is validated by designing 10 novel single-site mutants of a DNA polymerase, which wet-lab experiments show have enhanced activity at elevated temperatures, demonstrating practical utility for enzyme engineering.

Introduction

Enzymes are central to synthetic biology, but wild-type variants often suffer from low activity, poor stability, or weak binding, motivating enzyme engineering to tailor these properties. Deep learning models pretrained on large protein sequence corpora have emerged as powerful tools for predicting mutation effects, yet existing methods typically capture only one or two modalities: sequence, structure, or evolutionary information from multiple sequence alignments. No prior approach integrates all three, and those that incorporate homology data either require extra training or lack plug-and-play flexibility, leaving learned representations potentially incomplete. The authors address this gap with VENUSREM, a pretrained protein language model featuring a retrieval-enhanced module that fuses sequence and structure tokens through disentangled cross-attention layers and seamlessly injects evolutionary signals via an alignment tokenizer without additional training. This unified representation consistently outperforms all models on the ProteinGym benchmark across 217 assays and reliably identifies beneficial single- and multisite mutations in both in silico and wet-lab experiments, demonstrating strong potential for practical enzyme engineering.

Dataset

Here is a concise dataset description based on the provided paper paragraphs.

Sequence tokens

• The vocabulary contains 20 standard amino acids and 5 special tokens (<pad>, <cls>, <eos>, <unk>, <mask>).
• The pretrained language model is trained on data from Foldseek, which clusters the AlphaFold2 database and retains the original PDB structures and amino acid sequences.
• For inference, only the wild-type protein sequence and structure are used; no mutation masking or position-wise mutating is applied.

Structure tokens

• A local structure codebook is built from 4,735,677 local structure graphs extracted from CATH43-S40.
• For each residue, a local neighborhood is defined by selecting up to 40 neighboring residues within a 10 Å spatial radius. An undirected graph connects residues with a distance ≤10 Å.
• The graphs are processed by a six-layer geometric vector perceptron (GVP) autoencoder, trained with a denoising objective (Gaussian noise on Cα coordinates and Brownian motion).
• A mean pooling operation produces a 256-dimensional vector, which is then mapped via K-means clustering to a 2048-dimensional discrete latent space (implicit structure vocabulary).

Homologous sequences Three retrieval strategies supply evolutionary information for each protein. For the 217 ProteinGym assays, the first strategy uses pre-downloaded MSAs; the other two are computed on the fly.

• EVcouplings alignment: Jackhmmer searches on UniRef100 with five iterations and bit scores from 0.1 to 0.9 (nine search groups). The .a2m file with the most significant evolutionary couplings is selected. Preprocessing converts lowercase letters to uppercase and replaces special characters with the <pad> token.
• ColabFold alignment: ColabFold queries a local PDB database and outputs .a3m alignments, which are then realigned with the reformat.pl script for compatibility.
• Foldseek alignment: Structural alignment is performed via the Foldseek API against databases (afdb50, afdb-proteome, cath50, pdb100). Gaps are removed from the query sequence, and the target alignment is padded to match the query length.

Evaluation datasets

• ProteinGym (substitution subset): 217 high-throughput assays with more than 2.5 million mutations.
• VHH antibody: 31 mutants (1–4 site mutations) from Kang et al. (2025), evaluated for binding affinity and alkali resistance. The structure is predicted by AlphaFold3, and homologous sequences are retrieved from UniRef100.

How the data is used

• The sequence tokenizer is pretrained on the Foldseek clustered data.
• The structure tokenizer is trained on the 4.7M local structure graphs from CATH.
• The model takes wild-type sequence tokens, structure tokens, and processed homologous alignments from the three retrieval pipelines. It is then evaluated on the ProteinGym and VHH antibody benchmarks.

Method

The VENUSREM framework is a zero-shot predictor for mutation effects, accepting a wild-type protein’s sequence, structure, and homologous sequence information as inputs. For a given protein of length L, the sequence is directly tokenized into one-hot vectors over the 20 standard amino acids plus five special tokens (, , , , ). Local structure is encoded by a discrete tokenization pipeline: for each residue, a graph is built from up to 40 spatial neighbors within a 10 Å radius, and a pretrained geometric vector perceptron (GVP) autoencoder maps each graph to a 256-dimensional continuous vector through mean pooling. These continuous vectors are then quantized via K‑means clustering into a 2048‑dimensional codebook, yielding a discrete structural token per residue. The GVP autoencoder is trained with a denoising objective, perturbing Cα coordinates with 3D Gaussian noise and Brownian motion.

The sequence and structure tokens are fed into a BERT‑style protein language model (PLM) that computes a native representation using disentangled multi‑head cross‑attention. For positions i and j, the attention score is constructed from the sequence token embeddings R_i, R_j, the structure token embeddings S_i, S_j, and the relative positional encoding P_{ij}. The initial attention is expressed as:

$$\mathrm{Attn}(i, j) = \{\boldsymbol{R}_i, \boldsymbol{S}_i, \boldsymbol{P}_{ij}\} \times \{\boldsymbol{R}_j, \boldsymbol{S}_j, \boldsymbol{P}_{ji}\}^\top,$$Attn(i,j)={Ri​,Si​,Pij​}×{Rj​,Sj​,Pji​}⊤,

which, after dropping terms that contain no residue information, simplifies to five AA‑relevant components:

$$\begin{aligned} \mathrm{Attn}(i, j) &= \boldsymbol{R}_i \boldsymbol{R}_j^\top + \boldsymbol{R}_i \boldsymbol{S}_j^\top + \boldsymbol{R}_i \boldsymbol{P}_{ji}^\top + \boldsymbol{S}_i \boldsymbol{R}_j^\top + \boldsymbol{P}_{ij} \boldsymbol{R}_j^\top \\ &= \boldsymbol{Q}_i^R (\boldsymbol{K}_j^R)^\top + \boldsymbol{Q}_i^R (\boldsymbol{K}_j^S)^\top + \boldsymbol{Q}_i^R (\boldsymbol{K}_{ji}^P)^\top \\ &\quad + \boldsymbol{Q}_i^S (\boldsymbol{K}_j^R)^\top + \boldsymbol{Q}_{ij}^P (\boldsymbol{K}_j^R)^\top. \end{aligned}$$Attn(i,j)​=Ri​Rj⊤​+Ri​Sj⊤​+Ri​Pji⊤​+Si​Rj⊤​+Pij​Rj⊤​=QiR​(KjR​)⊤+QiR​(KjS​)⊤+QiR​(KjiP​)⊤+QiS​(KjR​)⊤+QijP​(KjR​)⊤.​

The attention scores are scaled by $1/\sqrt{5d}$1/5d​ (with $d$d the dimension of $\boldsymbol{Q}^R$QR), normalized with softmax, and used to update the hidden representation:

$$\boldsymbol{R}_o = \sigma\!\left(\frac{\boldsymbol{H}_{\mathrm{attn}}}{\sqrt{5d}}\right) \boldsymbol{V}^R.$$Ro​=σ(5d​Hattn​​)VR.

The PLM is pretrained on a large corpus of protein structures, so it produces a native logits tensor $\boldsymbol{O}_{iv}^{\mathrm{native}}$Oivnative​ for each position i and vocabulary index v.

In parallel, homologous sequences for the query protein are retrieved using one of three strategies: EVcouplings‑based MSAs, ColabFold‑automated MSA with realignment, or Foldseek structural alignment. The aligned sequences are assembled into a multiple sequence alignment (MSA) of N homologous sequences. For each residue position i, a frequency vector is computed by counting the occurrences of each amino acid type across the N sequences, normalized to form a counting matrix $\mathbf{C}$C:

$$\mathbf{C}_{iv} = \frac{\sum_{n=1}^{N} \mathcal{I}(\mathbf{A}_{ni} = v)}{\sum_{v=1}^{V} \sum_{n=1}^{N} \mathcal{I}(\mathbf{A}_{ni} = v)},$$Civ​=∑v=1V​∑n=1N​I(Ani​=v)∑n=1N​I(Ani​=v)​,

where $V=25$V=25 and $\mathcal{I}$I is the indicator function. The evolutionary logits are then obtained by applying a softmax followed by a log transformation:

$$\mathbf{O}_{iv}^{\mathrm{evo}} = \log\!\left(\frac{\exp(\mathbf{C}_{iv})}{\sum_{v=1}^{V} \exp(\mathbf{C}_{iv})}\right).$$Oivevo​=log(∑v=1V​exp(Civ​)exp(Civ​)​).

The final output logits are a weighted combination of native and evolutionary contributions, controlled by a parameter $\alpha \in [0,1]$α∈[0,1]:

$$\boldsymbol{O}_{iv}^{\mathrm{out}} = (1 - \alpha) \cdot \boldsymbol{O}_{iv}^{\mathrm{native}} + \alpha \cdot \boldsymbol{O}_{iv}^{\mathrm{evo}}.$$Oivout​=(1−α)⋅Oivnative​+α⋅Oivevo​.

For a given mutant with a set of mutated positions $\mathcal{T}$T, where each position $t$t changes the wild‑type residue $v$v to $v'$v′, the zero‑shot fitness score is the sum of the logit differences at all mutated sites:

$$F_x = \sum_{t \in \mathcal{T}} (\boldsymbol{O}_{t v'}^{\mathrm{out}} - \boldsymbol{O}_{t v}^{\mathrm{out}}).$$Fx​=t∈T∑​(Otv′out​−Otvout​).

This scoring scheme captures the predicted change in fitness relative to the wild type without requiring any mutation‑specific training, relying solely on the integrated sequence, structure, and evolutionary signals.

Experiment

The evaluation of VENUSREM spans high-throughput benchmarking on ProteinGym, post-hoc analysis of VHH antibody mutants, and prospective wet-lab engineering of phi29 DNA polymerase. Integrating sequence, structure, and evolutionary information, VENUSREM achieves top leaderboard performance and robust fitness predictions across diverse property types. It reliably captures multi-site mutation effects on antibody binding and alkali resistance, and in a wet-lab setting, it designed enzyme variants that simultaneously improved activity and thermostability, confirming its practical utility for protein engineering.

VenusREM achieves the highest overall average Spearman correlation on the ProteinGym benchmark, outperforming all other models. It ranks first in activity, binding, and expression predictions, while achieving second in stability and third in organismal fitness. By integrating sequence, structure, and evolutionary information, it surpasses models that rely on only a subset of these modalities. VenusREM attains the best overall average Spearman correlation, noticeably ahead of the runner-up ProSST. It leads in activity, binding, and expression predictions, and is a close second in stability behind ProSST. Integration of sequence, structure, and evolutionary features enables VenusREM to outperform models that use only one or two of these inputs. In organismal fitness, evolution-aware models PoET and TranceptEVE take the top two spots, with VenusREM placing third.

On the ProteinGym benchmark, VenusREM achieves the highest overall average Spearman correlation, leading in activity, binding, and expression predictions, while ranking second in stability and third in organismal fitness. Its integration of sequence, structure, and evolutionary features allows it to outperform models that rely on only one or two of these modalities. The results highlight that combining multiple data types yields superior performance across diverse protein function prediction tasks.