Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
Authors: Ruben Weitzman
Peter Mørch Groth
Lood Van Niekerk
Aoi Otani
Yarin Gal
Debora Susan Marks
Pascal Notin
Venue: ICML 2025 poster (2025)
Keywords: Computational Biology
Primary area: applications->health_medicine

Abstract: Retrieving homologous protein sequences is essential for a broad range of protein modeling tasks such as fitness prediction, protein design, structure modeling, and protein-protein interactions. Traditional workflows have relied on a two-step process: first retrieving homologs via Multiple Sequence Alignments (MSA), then training mod- els on one or more of these alignments. However, MSA-based retrieval is computationally expensive, struggles with highly divergent sequences or complex insertions & deletions patterns, and operates independently of the downstream modeling objective. We introduce Protriever, an end-to-end differentiable framework that learns to retrieve relevant homologs while simultaneously training for the target task. When applied to protein fitness prediction, Protriever achieves state-of-the-art performance compared to sequence-based models that rely on MSA-based homolog retrieval, while being two orders of magnitude faster through efficient vector search. Protriever is both architecture and task-agnostic, and can flexibly adapt to different retrieval strategies and protein databases at inference time – offering a scalable alternative to alignment-centric approaches.

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
Ruben Weitzman 12
Peter Mørch Groth 34 Lood Van Niekerk s
Aoi Otani?
Yarin Gal Debora S. Marks
Pascal Notin
Abstract
of function shapes the landscape of permissible mutations (Göbel et al.. 1994). Understanding these mutational land- scapes is fundamental to both basic biological research and the engineering of novel proteins for applications in ther- apeutics, new materials, and sustainability. Homologous proteins sequences sharing a common evolutionary ori- gin - are particularly valuable for modeling because they often exhibit similar structural and functional properties that reveal fundamental constraints on sequence variation. Con- sequently, leveraging homology has become essential for range of protein modeling tasks, from predicting the effects of mutations on disease risk (Frazer et al., 2021), ), to guiding protein design campaigns (Russ et al., 2020; Notin et al., 2024), or inferring tertiary structure from sequence (Jumper et al., 2021).
Retrieving homologous protein sequences is es- sential for a broad range of protein modeling tasks such as fitness prediction, protein design, struc- modeling, and protein-protein interactions. Traditional workflows have relied two-step process: first retrieving homologs via Multiple Sequence Alignments (MSA), then training mod- els on one or more of these alignments. However, MSA-based retrieval computationally expen- sive, struggles with highly divergent sequences or complex insertions & deletions patterns, and op- erates independently of the downstream modeling objective. We introduce Protriever, end-to-end differentiable framework that learns to retrieve relevant homologs while simultaneously training for the target task. When applied to protein fitness prediction, Protriever achieves state-of-the art per- formance compared to sequence-basee models that rely on MSA-based homolog retrieval, while being two orders of magnitude faster through effi- cient vector search. Protriever is both architecture- and task- aask-agnottic, and can flexibly adapt to dif- ferent retricval strategies and protein databases at inference time- - offering a scalable alternative to aigmment-centrii approaches.
Traditional protein modeling workflows follow two-stage process: first retrieving homologs via Multiple Sequence more of
Alignments (MSAs), then training models these alignments. This approach has yielded family-specifh models that effectively capture evolutionary constraints for particular protein families (Krogh, 1998; Hopf et al., 2014; Frazer al., 2021). However, this paradigm suffers from several fundamental limitations. First, MSA-base retrieval often misses distantly related sequences that could provide valuable evolutionary context but fall below alignment sig- nificance thresholds. Second, sequences containing large insertions, deletions, or structural rearrangements remain difficult to align reliably, despite potential functional re- latedness (Riley et al., 2023). Third, retrieval operates in- dependently of downstream modeling objectives, relying alignment heuristics rather than principled, data-driven approaches to identify the most informative homologs for specific tasks. Additionally, analyzing new protein families requires constructing separate MSAs and training distinct models, making the approach computationally demanding and poorly suited for large-scale applications.
1. Introduction
Proteins have evolved over billions of years under strict con- straints and evolutionary pressure, where the preservation
Department of Computer Science, University of Ox- ford *Department of Systems Biology, Harvard Med- Computer Science, Uni-
ical School 'Department versity of Copenhagen *Enzyme research, Novonesis 'Ginkoo Bioworks. Correspondence to: Ruben Weitzman huubee.weitzmant wcw..x.............. Debora Marks <debb bic@hm hms hawvard..dduu. Yarin Gal aaarinns..ox.ac.uk Pascal Notin <pascal notin hhmmshan harvard.edu>..
Recent advances in large-scale protein language models (pLMs) have yielded flexible, alignment- approaches that leverage the vast diversity of known protein sequences (Einaggar et al., 2021; Lin et al., 2023; Nijkamp et al., 2023). However, single-sequence models often underper- form family-specifh methods for variant effect prediction,
Proceedings of the nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
S
Q
Query sequence
Retrieved sequences
Reconstructed query
M I IQH..
Retriever
M
S
Q
H
Reader
M Q
M
T
L
D
M
T
S
Query sequence M
Retriever encoder
Query embedding
Complete retrieved set
Filtering sampling
Filtered retrieved set
K A
Reader loss
d.
dec
BETR
loggrartaileee d.)
Protein database (e.g., UniRef50)
Database embeddings
Retriever loss
Vector similarity search
Laata
Conditioning of homologous sequences
Reader
Autoregressive decoding of
log
Σ Paal (inarnddc
4.
Index
d:
Hybrid loss
= CRETR +oCe
LM
Figure 1. Protriever. The Protriever framework is composed of three parts: learned retriever, an index, and reader. The two neural
networks work together to produce conditional sequence likelihood. The retriever selects set of sequences Dx {d4 l..... . to be passed on to the reader using vector similarity search between the embedded query sequence and retrieval index. This set of sequences then passed the reader which learns to reconstruct the query from the filtered conditioning set Dxy. During training, the reader calculates the relevance score of each document PLM (q dę) which are then used to train the retriever.
particularly when confronted with rare or highly specialized proteins (Notin et al., 2023). Hybrid solutions therefore in- tegrate family-specific context through several approaches. Some methods fine-tune pretrained models on homologous sequences for specific protein families (Alley et al., 2019), while others learn distributions over fixed sets retrieved sequences during training (Rao et al., 2021; Truong Jr & Bepler, 2023). A third strategy performs retrieval at infer- ence time, integrating family-specific statistics with learned representations across families (Notin et al., 2022). Despite this progress, existing retrieval frameworks remain static, where retrieval depends on fixed notions of sequence simi- larity without allowing models to refine or backpropagate through retrieval choices.
triever pretraining strategies such as Dense Passage Re- trieval (DPR), and different sampling strategies to opti- mize fitness prediction performance ($ 3.2 and $ 4.1);
We leverage and adapt several architectural speedups for fast vector search, including inverted file indexes and product quantization ($ 3.3);
We demonstrate Protriever achieves state-of-the-ard per- formance among sequence- sequence-based based models on the Pro- teinGym benchmarks, while being orders of magni- tude faster to retrieve homologs than standard MSA approaches including JacKHMMER, MMseqs2, and MMseqs2-GPU ($ and 5).
In this work, propose Protriever, an end-to-end differ- entiable protein homology search for fitness prediction. Our approach enables fast vector-based homology retrieval using learned dense representations and integrates these retrieved sequences to yield accurate zero-shot fitness predictions. Our contributions are as follows:
The integration of retrieval mechanisms directly into the training process addresses crucial gaps in traditional protein sequence modeling. First, it overcomes the fixed nature of homology sets and their limitations on adapting to newly discovered or less common sequences. Second, by oper- ating in dense representation space rather than requiring explicit alignments, Protriever can leverage distantly related structurally divergent sequences that are difficult to align reliably but may still provide valuable evolutionary con- text. This dual advantage allows the model to dynamically discover informative homologs while accessing a broader range evolutionary relationships than alignment based approaches.
We develop end-to-end differentiable retrieval and sequence modeling approach, enabling our joint archi- tecture to learn which homologous sequences are the most informative for the downstream task of interest (S 3.1);
We explore various hybrid losses for joint training, re-
2

---

Protriever: End-to-End Differentiable Protcin Homology Search for Fitness Prediction
2. Related work
2.3. Hybrid fitness prediction models
2.1. Alignment-based models
Hybrid approaches combine the generalization capabili- ties of pLMs with evolutionary specificity through three main strategies. The first involves fine-tuning pretrained models with the same self- objective on sub- set of family-specific sequences, typically retrieved with a MSA, as demonstrated by UniRep's ' evolutionary fine- tuning (evotuning) (Alley et al., 2019) or ESM-Iv's spiked fine -eeitinig (Meier et al., 2021). The second strategy trains models across sets of homologous sequences during training. MSA Transformer (Rao et al., 2021) learns distributions over entire MSAs using axial attention and masked lan- guage modeling, enabling rich cross-family representations. PoET (Truong Jr & Bepler, 2023) and ProtMamba (Sgar- bossa et al., 2024) eliminate alignment requirements by learning distributions over concatenated homologous se- quences. The third approach performs retrieval infer- ence time. Tranception (Notin et al., 2022) weights predic- tions from unconditional language model with position- specific frequencies from MSAs. while TranceptEVE fur- ther integrates predictions based on a model that learns de- pendencies across positions in the MSA. However, all these approaches rely on predefined homology sets determined by traditional sequence similarity heuristics, failing to leverage model-learned representations for optimal retrieval.
Alignment-buu models have been the cornerstone of pro- tein sequence analysis for decades. The traditional workflow involves searching large protein databases for homologs, constructing Multiple Sequence Alignments (MSAs), and fitting statistical models to capture evolutionary constraints within protein families. Early methods like PSSM (Position- Specific Scoring Matrices) and HMMs (Hidden Markov Models) extracted position-specific information from align- ments (Krogh, 1998). Subsequent approaches captured pair- wise co-evolutionary signals through energy-based mod- els (Hopf al., 2017), while more recent methods leverage variational autoencoders to model higher-order correlations between sequence positions (Riesselman et al., 2018; Frazer et al., 2021). However, alignment- based methods face fun- damental limitations that have persisted across generations of approaches. MSA construction algorithms like MUS- CLE (Edgar, 2004), ClustalW (Thompson et al., 1994), and more recent tools like JacKHMMER (Johnson et al., 2010) and MMseqs2 (Steinegger& & Söding, 2017) rely se- quence similarity heuristics that can fail for highly divergent sequences. These methods struggle with regions containing extensive insertions, deletions, or structural rearrangements, often producing gaps that obscure true evolutionary rela- tionships. Additionally, MSA-based -based models must be trained separately for each protein family, making them computa- tionally expensive and limiting their applicability to families with sparse sequence data. The fixed nature of these align- ments also means that as protein databases grow with newly sequenced genomes, existing models become obsolete and require complete MSA reconstruction and retraining.
2.4. Retrieval
Retrieval methods identify relevant database objects im- prove task performance. In NLP, retrieval has enhanced open-domain question answering (Robertson & Zaragoza, 2009; Grave et al., 2017; Wang et al.. 2018; Karpukhin et al., text generation (Lee et al., 2019). Advanced methods like REALM (Guu et al., 2020) and RAG (Lewis et al., 2020) demonstrated end-to-end differentiable training of retriever-reader architectures for retrieval- augmented generation. In protein modeling, MSA generation has been framed as retrieval, with some work ex- ploring differentiable MSA construction (Hong et al., 2021; Petti et al., 2023; Llinares-López et al., 2023). Recent ap- proaches include RAG (Li et al., 2024), which trains retrievers for improved MSA generation by using hierarchi- cal clustering to generate sequence identifiers that a causal language model learns to predict, and RSA (Ma al., 2024), which uses frozen ESM- lb encoders for sequence retrieval. However, no prior work has achieved end-to-end joint train- ing of retrieval and protein sequence modeling.
2020) and snowledge-groundee
2.2. Protein language models
Protein language models (pLMs) emerged as alignment- free alternatives, trained with self-supervised objectives on massive protein sequence databases. These models learn evolutionary constraints that generalize across families, par- ticularly benefiting small families with limited homologs. UniRep (Alley et al., 2019) pioncered cross-family pro- tein modeling using LSTM architectures. The field then adopted Transformer-based approaches, including autore- gressive models like ProGen (Madani et al., 2020), Tran- ception (Notin et al., 2022), and ProtGPT2 (Ferruz et al., 2022), and masked language models like ESM (Rives et al., 2021), PRoBERTa (Nambiar et al., 2020), and ProtTrans (EI- naggar al., 2021). Despite their flexibility, pLMs often underperform family-specific models for fitness prediction and require substantial computational resources to encode evolutionary knowledge across all protein families in their parameters.
3. Protriever
3.1. Overall architecture
The Protriever framework is comprised of three components: the Retriever model, the Index, and the Reader model (see
3

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
Figure 1). At high level, the query sequence is first pro- cessed by the retriever, which performs similarity search against a fixed index of sequence embeddings to identify rel- evant homologs. The reader model then performs the target task conditioning on the retrieved homologous sequences. During training, the target task is typically a self-supervised objective such as autoregressive decoding of the query se- quence. The reader learns which retrieved sequences pro- vide useful context for that task, providing gradient feedback to the retriever that adjusts sequence relationships in embed- ding space accordingly. The following subsections detail each component and the training procedure.
tional to the size of the corresponding UniRef90 cluster. As data augmentation, we randomly reverse query sequences during training from either N-terminus to C-terminus, or vice versa, to make retrieval agnostic to sequence order.
This pretraining phase establishes a foundation for the subse- quent end-to-end training of the full Protriever architecture, ensuring that the retriever can effectively identify relevant protein sequences from the UniRef50 index for downstream reconstruction tasks.
3.3. Index
The index consists of embeddings for all UniRef50 (Suzek et al., 2015) sequences (R: 62 million), computed using the (pretrained) retriever encoder and stored for efficient similar- ity search via Faiss (Johnson et al., 2021). As the retriever encoder is updated during training, the index embeddings can become stale relative to the evolving retriever represen- tations. Since updating the index is computationally expen- sive-requiring re-encoding all 62 million sequences-we update it 10 times throughout training (every 5k training steps). This approach balances maintaining reasonably cur- rent embeddings with computational efficiency.
3.2. Retriever module
Initialization We use transformer encoder architecture for the retriever, which we imitialize with ESM-2 (Lin et al., 2023) pre-trained weights (35M parameters). Average pool- ing is applied over the outputs of the last layer to obtain a 480-dimensional vector representation for each sequence. To then compute similarity between query q and other sequences in the index, we compute the cosine similarity s(d, q) between their corresponding embeddings.
Pretraining with dense passage retrieval Following the dense passage retrieval paradigm (DPR, Karpukhin et al. (2020)), we further pretrain the retriever encoder to learn effective protein sequence representations in a low- dimensional continuous space. The core objective is to create an embedding space where homologous protein se- quences have higher similarity scores than non-homologous ones. We construct training data from UniRef50 us- ing BLAST all-vs-all search to identify homologous se- quences for each query (see Appendix G). Let D
For efficient search at scale, we implement two key opti- mizations. First, we use an inverted file index (IVF), where the embedding datahase is partitioned using k-means clus- tering into KIvP clusters. At query time, we search only the nearest Pive clusters rather than the entire database. Sec- ond, we apply product quantization (PQ) to compress the stored vectors, reducing memory usage while maintaining retrieval accuracy. Due to the scale of UniRef50, the index is distributed across multiple GPUs, with queries processed independently on each partition and results aggregated to produce the final retrieval set. Detailed specifications are provided in Appendix A.
{(g: dtt,,..,dd d m, d.......,dd .....,,,d.... data consisting of M instances. Each instance contains one query sequence qı m, relevant (positive) homologous sequences d and irrelevant (negative) sequences d Given this data, we optimize the encoder parameters to max- imize the similarity s(d+, q) while minimizing s(d=,q) ,4) using the negative log likelihood:
) be the training
3.4. Reader module
We use PoET (Protein Evolutionary Transformer) (Truong & Bepler, 2023) as our reader backbone. PoET is designed to operate on sets of homologous sequences by processing them as concatenated sequences-of-sequeee -equennecssssssssseeeeeeeeeeeeee through spe- cialized decoder transformer architecture. Unlike alignment- based approaches, PoET learns to model evolutionary re- lationships without requiring explicit sequence alignment, making it well- suited for integration with our retrieval frame- work. The PoET architecture processes the query sequence concatenated with all retrieved homologous sequences si- multaneously, allowing the model to capture complex in- teractions between all positions across the entire set. This enables the reader to learn rich representations that incorpo- rate evolutionary context from the retrieved sequences when making predictions for the query.
m,
Cpeetrain
log (andts))) + En_ (7))
Negative examples are comprised of both random sequences and positive examples for other queries from the same mini- batch. During training, we sample UniRef50 clusters with weight inversely proportional to cluster size to avoid over- representing large clusters (see Figure G.1). We then replace each UniRef50 sequence, whether query or cluster member, with a UniRef1 00 sequence with weight inversely propor-
4

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
We initialize our reader with a pretrained PoET model that was trained on the same UniRef50 dataset used for our re- triever's DPR pretraining. This initialization provides strong foundation for subsequent end-to-ent training, as the model already understands how to process sets of homolo- gous sequences effectively. Our framework is general and can adapt to various model architectures and self-supervised objectives, as long as these architectures can condition their predictions sets of homologous sequences. We present architecture
in a sequence of amino acids given retrieved sequences as context. The standard autoregressive factorization is ex- tended by conditioning on a set of K retrieved sequences Dx = tott-o (d/z)):
P(z) = PRETR (Dx(z) II (E Dx).
results with an efficient sequence-to-sequee (Fusion-in-Decoder) in Appendix B, demonstrating the ver- satility of our approach.
Following Frazer al. (2021) and Notin et al. (2022), we evaluate the fitness of mutated protein sequence æ- via its log- ligeeiidlhoodd ratio with respect to the wild-type sequence xw:
3.5. Protriever training
P (x:
We evaluate different loss functions for end-to-end train- ing of the retriever. Our approach leverages the language model's 's performance to guide retriever training: mologous protein proves valuable for the reader's sequence modeling task, the retriever is encouraged to rank closer to the query sequence in embedding space.
Fg = log P (x:
ho-
In practice, both mut and x are close in sequence space (e.g., differ by handful of mutated positions), they will share the same conditioning set Dx. - In that case, the retriever probabilities cancel out, simplifying the fitness score to:
The relevance score of protein sequence d to query sequence q is computed as:
PLM (a Dx) F2 = log PIM (w (Dx)
PRETR(d 9) = exp(s(d, q)/0) E exp (s (dx; 9) /0)
(2)
To apply this scoring methodology, first build an index of all protein sequences in our database. At inference time, we use the trained retriever from Protriever to encode all 62 million UniRe150 sequences. This process is parallelized across GPUs and uses FlashAttention (Dao et al., 2022) to enable large batch sizes, completing in approximately 30 minutes on four A100 GPUs. We then construct a Faiss index for fast similarity search in 3-4 minutes across four GPUs (see details in Appendix A).
where the sum is over Dx= = =d+ top-K retrieved sequences, approximating the true relevance score over the entire index while maintaining computational tractability.
For our main experiments, we use the End-to-end train- ing of Multi-Documen Reader and Retriever (EMDR) loss (Sachan et al., 2021), which treats retrieved sequences latent variables. Given query q and the set Dx of top-K retrieved sequences, the loss is:
Given a query sequence. we retrieve relevant homologs through a procedure inspired by re-ranking approaches in RAG (Izacard et al., 2022) and diversity maximization strate- gies in alignment-basee based protein modeling (Hopf et al., 2017: Rao al., 2021). Re-ranking allows
LEMOR log Σ PLM (q d&) PRETR (d& ②
refine initial retrieval results by considering additional factors beyond similarity scores-in our case, diversity maximization to en- sures we capture a broad range of evolutionary relationships. Following PoET (Truong Jr & Bepler, 2023), we enhance fitness prediction performance by ensembling F, estimates across multiple conditioning sets D that vary in size and composition.
(3)
During optimization, we apply a stop-gradient operator to PLM. ensuring updates are limited to the retriever parame- ters. In earlier stages of model development, we also experi- mented with alternative loss functions, including Perplexity Distillation (PDist) and Leave-One One-Out Perplexity Distilla- tion (LOOP) (see Appendix D).
More specifically, we search the UniRef50 index for homol- ogous sequences and pull corresponding sequences from UniRef100. We sample based on effective cluster size and distance to query, evaluating five parameter combinations, along with three conditioning set sizes (6k, 12k, and 24k tokens), yielding an ensembling strategy across 15 forward
4. Fitness prediction with Protriever
4.1. Fitness scoring methodology
The reader model of Protriever is trained with a condi- tional autoregressive objective, predicting the next token
5

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
Table Zero-shot performance on substitution DMS of ProteinGym benchmark. Reported metrics are Spearman rank correlation, AUC, MCC, top recall, and NDCG. Models are classified according to if they take as input MSAs (alignment based and Hybrid) not (unconditional pLMs and Protriever)
Model type
Model name
Spearman
AUC
MCC
Recall
NDCG
Site independent
0.359
0.696
0.287
0.201
0.748
Alignment-
GEMME
0.459
0.749
0.353
0,211
0.777
based
EVE
0.439
0.742
0.342
0.229
0.782
ESM-1v
0.407
0.724
0.321
0.210
0.749
Unconditional
ProGen2
0.391
0.717
0.306
0.198
0.767
pL.
ESM2
0.405
0.726
0,322
0.213
0.764
MSA Transformer
0.432
0.737
0.341
0,223
0.777
Tranception L
0.434
0.741
0.341
0.220
0.779
Hybrid
TranceptEVE I.
0.458
0.754
0.356
0.229
0.786
PoET
0.470
0.759
0.368
0.226
0.784
Protriever
Protriever
0.479
0.762
0.374
0.229
0.788
passes with different configurations. Additional details the sampling procedure provided in Appendix E.
ble Our approach reaches Spearman correlation of 0.479 across all assays (aggregated by protein ID), out- performing the previous best sequence-only model PoET (0.470) and all other baselines. Protriever achieves the best performance across all metrics: AUC (0.762), MCC (0.374), NDCG (0.788), and top-K recall (0.229).
4.2. Experimental setup
We evaluate Protriever on the substitution benchmark of ProteinGym (Notin et al., 2023), containing 217 deep mu- tational scanning (DMS) experiments that probe the nat- ural function of protein variants. DMS experiments sys- tematically measure the functional effects of individual amino acid substitutions across a protein sequence, provid- ing comprehensive fitness landscapes for specific proteins. Consequently, to perform well this benchmark, models must capture nuanced understanding of the biochemical constraints for the corresponding proteins as they must be able to detect subtle effects resulting from minor sequence changes.
These results demonstrate that end-to-end differentiable re- trieval effectively identify and leverage homologous sequences for fitness prediction, achieving competitive per- formance with hybrid models that rely on traditional MSA construction while maintaining the computational efficiency and flexibility learned retrieval. Additional results seg- mented by MSA depth and taxonomic groups are provided in Tables C.1 and C.2, showing that Protriever performs consistently well across different protein families and evo- lutionary contexts, with partieularly strong performance prokaryotes and viruses.
Performance evaluated using multiple metrics: Spearman correlation, AUC, and MCC capture broad performance across the full assay (important for general mutation ef- feet prediction), while NDCG and top-K recall focus the top end of the measured phenotypes (and thus most important for protein design applications). To evaluate the zero-shot capability our retrieval framework, we score all sequences within each DMS using the same conditioning sets with homologous sequences described above. Addi- tionally, score sequences both directions (N-terminus to C-terminus and vice versa), strategy shown to improve predictive performance (Notin et al., 2022).
5. Discussion
5.1. Model training ablations
We conduct an ablation study to understand the individ- ual contributions DPR initialization and end-to-end joint training in our Protriever framework. Table 2 compares four experimental configurations in terms of Spearman correla- tion on the ProteinGyym DMS substitution benchmark across different MSA depth regimes.
Using ESM embeddings directly as retriever with our in- ference filtering and sampling scheme achieves Spearman correlation of 0.432. This baseline demonstrates that ESM embeddings capture meaningful sequence homology rela- tionships despite not being explicitly trained for retrieval
4.3. Results
Protriever achieves state-of-the-ar performance among sequence-based models ProteinGym, shown in Ta-
6

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
Protriever (Faiss)
MMseqs2-GPU
MMseqs2 (k-mer)
JacKHMMER
0.08
10
0.07
0.06
Time per query (s)
10
0.05
Density
0.04
10
0.03
0.02
10
0.01
0.00
10 100
1000
20
40 60
100
Number of queries
Sequence identity between queries and retrieved sequences
Figure 2. Retrieval speed and quality. Left: Retrieval time per query sequence (mean and standard error) at different query sizes using embedding similarity scarch, MMscqs2, and GPU-accelerated MMseqs2. Benchmarking details and tabulated values can be found in Appendix A.3 and Table A.2, respectively. Right: Distribution of sequence identities between queries and retrieved sequences for Protriever. MMSeqs2, and JacHHMMER, using the ten ProteinGym validation sets as queries. The retrieved sequences have not been filtered and weighed by number of retrieved sequences per method/query pair. Retrieval methods sensitive hyperparameters. We provide additional results for JMMRR with bit scores 0,1, 1,1, 0.5 (shown ahove), and 0.9 in Appendix A.4.
tasks, and that our retrieval framework can effectively lever- age these relationships for fitness prediction.
ilarity search. The MSA sequences used by methods like EVE (Frazer et al., 2021) rely on sensitive retrieval tools such as JackHMMER (Johnson et al., 2010), which can take hours or days to complete for a given query sequence. Other retrieval routines such as BLAST (Altschul al., 1990) MMseqs2 (Steinegger & Söding, 2017) can offer speed improvements but at the of lower retrieval pre- cision. Recently, MMseqs2-GPU (Kallenborn al., 2025) further reduced search times to hundreds of milliseconds with marginal impact on retrieval accuracy. In contrast, our approach leverages fast vector similarity search using the Faiss library (Douze et al., 2024). Given the pre-trained index, retrieval rapid. lightweight, and scalable as demon- strated in Appendix A.
Replacing the ESM retriever with DPR improves perfor- mance 0.440, representing substantial gain. This im- provement highlights the advantage of contrastive leamiing for sequence retrieval, DPR is specifically trained known distant sequence homologs. Training the retriever and reader jointly (Protriever without DPR initialization) achieves 0.466, outperforming previous baselines. This demonstrates that joint optimization allows the retriever to learn representations better suited for the downstream fitness prediction task, with the reader model distilling informa- tion about homolog usefulness for sequence reconstruction across protein families.
To compare inference-time retrieval speed, we benchmark Protriever, MMseqs2 (Steinegger & Söding, 2017), and GPU-accelerated MMseqs2 (Kallenborn et al., 2025). The results are visualized the left-hand side of Figure and the benchmarking methodology is described Ap- pendix A.3. Our approach is two orders of magnitude faster than MMseqs2-GPU enabling significant efficiency gains for large-scale applications such as proteome-wide predic- tions.
Our full Protriever model, which combines DPR initial- ization with end-to-end joint training, achieves the highest performance at 0.479 Spearman correlation, validating our hypothesis that both components are essential. Analyzing performance by MSA depth reveals that both DPR initial- ization and joint training provide their largest benefits for low-depth MSAs, followed by medium-depth, then high- depth MSAs. This pattern suggests that learned retrieval is most valuable when traditional alignment-basee methods struggle due to limited evolutionary information.
We benchmark four methods to compare inference time retrieval speed and downstream performance: Protriever, MMseqs2, aaPelecclee MMseqs2, and JacKHMMER (Potter al., 2018). We search over UniRef50 clusters and apply similar inference scheme Protriever from Sec- tion 4.1. We apply the same 15% sequence identity cluster
5.2. Inference-time speedups
A key advantage of Protriever is the substantial speedup gained by replacing MSA-based retrieval with vector sim-

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
Table Protriever model performance on substitution DMS benchmark, Average Spearman's 's rank correlation between model scores and experimental measurements by MSA depth for different retrieval architectures.
Experiment
End-to-End
DPR
Spearman by
MSA depth
Low
Medium
High
Average
Frozen ESM
0.368
0.439
0.485
0.432
Frozen DPR
0.403
0.452
0.484
0.440
Protriever w/o DPR Protriever
V
0.461
0,476
0.508
0.466
0.464
0.498
0.512
0.479
filter as for Protriever, and sample according to combination of cluster sizes and distance the query, where distance is given by sequence similarity instead of embedding distance (details in Appendix E). PoET serves the reader model for all four configurations, with end-to-end training applied for Protriever.
different numbers of sequences with JacKHMMER often re- turning comparatively few high-identity sequences, whereas MMseqs2 and Protriever retrieve larger numbers of more diverse sequences as shown in Figure 2, which partly ex- plains HMMH'S sequence identity distribution. The distributions after applying our standard 15% sequence iden- tity filter (used during inference) are shown in Figure A.5. To further explore JacKHMMER's consitivity-perfommee trade-offss discussed above, we include results with vary- ing sensitivity settings both before and after filtering in Appendix A.4.
Spearman's 's rank correlation and average per-query retrieval time for the ProteinGym substitution benchmark are shown in Table The rapid retrieval of Protriever comes with- out performance loss for fitness prediction and even leads to gains through the joint retriever-reader framework. The CPU and GPU versions of MMseqs2 show nearly identical downstream performance as their hyperparameters were set to achieve similar sensitivity (see Appendix A.3 for details). JacKHMMER, despite being a higher sensitivity search tool, achieves slightly lower Spearman correlation, consistent with observations in the PoET paper (Truong Jr & Bepler, 2023). This may reflect how traditional search methods make retrieval decisions independently of the downstream task and reader model architecture, potentially leading to suboptimal sequence selection. This further underscores the importance of joint training to align retrieval with down- stream task modeling.
Interestingly, despite JackHMMER retrieving sequences with higher evolutionary similarity, it does not achieve su- perior fitness prediction performance. This suggests that the relationship between sequence similarity and utility for fitness prediction is not straightforward. and that learned retrieval representations may capture different aspects of protein relationships that are valuable for downstream tasks.
5.4. Portability and Modularity
Our vector-based retrieval framework introduces a paradigm shift in how homology search is performed and deployed. Unlike traditional MSA-based methods that require compu- tationally expensive sequence alignment at inference time, Protriever separates the indexing phase from the retrieval phase, enabling unprecedented portability and scalability.
5.3. Qualitative analysis of retrieved sequences
While we have demonstrated significant speedup in infer- ence time retrieval and improved downstream predictive performance, we have yet to examine the homology charac- teristics of the retrieved sequences. For ten validation sets (see Appendix F), we use wild-type sequences as queries and retrieve homologs using the same four methods as above. We compute sequence identity between queries and retrieved sequences and visualize the distributions on the right-hand side of Figure 2.
The core advantage lies in the transferable vector index. Once protein sequences are encoded in embedding space, the resulting index is lightweight, portable, and can be easily distributed or shared across different computational environ- ments. Our pre-computed index can be transferred as a com- pact file and immediately used for rapid retrieval without any additional preprocessing (see Appendix A, IVFPQ 6x8 in- dexing of UniRef50 uses 12.6GB of memory).
We observe that the distributions for MMseqs2 and Pro- triever searches have considerable overlap, while JackHM- MER results show substantially higher similarities to the queries. However, JackHMMER's 's more sensitive results come at significant computational cost as shown in Table 3. should also be noted that the four shown methods retrieve
This architecture enables dynamic database management without retraining. New sequences can be encoded and added to the index incrementally, while outdated or erro- neous entries can be removed instantly. This flexibility is particularly valuable for constantly evolving sequence datahases driven by advances in sequencing technology,
8

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
Table Retrieval methods performance comparison. Comparison of retrieval methods on ProteinGym substitution benchmark with retrieval timing. Inference is conducted using the scheme described in Section 4.1 with retrieval on UniRef50. Performance for MMseqs2 therefore different than as reported in original paper or Table. The reader model for all four entries is PoET, which for Protriever trained end-to-end using our framework.
Retrieval method
Spearman by MSA depth (f)
Retrieval time (s) (4)
Low
Medium
High
Average
Protriever
0.464
0.498
0.512
0.479
0.0046
MMseqs2 (k-mer)
0.455
0.472
0.489
0.463
16.860
MMseqs2- UMM
0.454
0.470
0.491
0.462
0.613
JacKHMMER
0.442
0.471
0.493
0.459
2501
or when working with proprictary sequences that cannot be publicly shared but can be incorporated into local in- dices. This approach contrasts with standard protein lan- guage models that encode sequence database information in their weights, making it difficult to leverage the information from newly added sequences without additional training.
This end-to-end optimization enables the model to discover functionally relevant evolutionary relationships that may be missed by alignment- approaches. Our modular frame- work design allows the trained retriever be paired with different reader architectures and adapted to various tasks beyond fitness prediction. This flexibility, combined with the computational efficiency of vector-based search, makes our approach broadly applicable across protein modeling applications. Additionally, our retrieval-based approach of- fers enhanced interpretability by allowing analysis of the sequences the model selects for conditioning, compared to traditional protein language models.
The framework also supports domain-specific specializa- tion. The same trained retriever can be applied to focused databases like GISAID (Shu & McCauley, 2017) for viral sequences or proprictary datasets for industrial applications, simply by re-indexing the target database without model retraining. This allows organizations to leverage the trained model while maintaining data privacy and domain relevance.
Future work will focus on scaling to larger databases and ex- ploring how the sequences retrieved by our learned approach compare to those selected by traditional MSA methods, be- yond sequence identity, providing deeper insights into the evolutionary relationships that drive protein function.
The framework is also model-agnostic. The retrieval frame- work can be integrated with various reader architectures, from encoder-decoder models (Appendix B) to decoder-only architectures. The retrieval encoder itself can be substituted with structure-aware models for tasks where structural in- formation is crucial, where structural homologs for training could be generated with TM-Vec (Hamamsy et al., 2024).
A Acknowledgements
We thank the members of the OATML group and Marks Lab for helpful discussions when writing this manuscript. R. W. is supported by the EPSRC Centre for Doctoral Train- ing in Health Data Science (EP/S02428X/1), P.M.G. was supported by Innovation Fund Denmark (1044-00158A), Y.G. holds Turing AI Fellowship (Phase 1) at the Alan Turing Institute, which is supported by EPSRC grant ref- erence V030302/1. D.S.M. and P.N. were supported by a Chan Zuckerberg Initiative Award (Neurodegeneration Challenge Network, CZ12018-191853) This research was also supported by a Dean's Innovation Award for the Use of Artificial Intelligence in Research from Harvard Medical School.
This modularity extends to training objectives. While we focus on autoregressive sequence reconstruction, the end-to- end differentiable framework can be adapted for other self- supervised tasks such as masked language modeling on the query sequence, supervised tasks such as property (c.g., thermostability, solubility) or tertiary structure predictions.
6. Conclusion
We have introduced Protriever, an end-to-end differentiable protein homology search framework that achieves state-of- the-art performance among sequence-based methods for protein fitness prediction while delivering computational speedups of two orders of magnitude over existing meth- ods for homology search. The key innovation lies in joint training of retrieval and prediction components, allowing the retriever to learn which homologous sequences are most informative for the specific downstream task rather than relying traditional task inseinndependent retrieval routines.
Impact statement
The advancement of protein design technologies promises transformative applications across many domains such as novel therapeutics, new materials, and sustainability. How- ever, as these technologies evolve, they raise important dual-

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
use concerns regarding their potential applications. While protein fitness and design models can accelerate therapeutic discovery, the same underlying capabilities could potentially be misused for harmful purposes, including the design of biological weapons. A notable example of this risk is how a model trained to minimize protein toxicity for drug develop- ment could theoretically modified to maximize toxicity instead (Urbina al.. 2022).
accuracy and high throughput. Nucleic Acids Research, 32(5):
1792-1797, 2004, doi: 10. 1093/nar/gkh3400
Elnaggar, A., Heinzinger, M., Dallago, C., Rehawi, G., Yu. W., Jones, L... Gibbs, T. Fehér, T. B., Angerer, C., Steinegger. M.. Bhowmik, D., and Rost, B. ProfTrans: Towards Cracking the Language of Lifes Code Through Self-Supervised Deep Learning and High Performance Computing. IEEE transactions on pattern analysis and machine intelligence. PP, 2021.
Ferruz, N., Schmidt, S., and Höcker, B. ProtGPT2 is deep unsupervised language model for protein design. Nature Com- munications, 13, 2022.
Our retrieval- approach offers potential advantages for addressing some of these concerns. The dynamic na- ture of RAG systems means we can extend the condi- tioning database without retraining, separating the model from the knowledge base. This architecture is particu- larly valuable for proprietary or sensitive biological se- orgnncss-orannization can use the pretrained Protriever while maintaining their sequences in private, access- controlled databases rather than incorporating them into model weights during training. Such flexibility allows or- ganizations to leverage the model's capabilities while main- taining strict control over their sequence data, which is especially important in biotechnology applications where data access may be restricted due to intellectual property concerns or biosafety considerations.
Frazer, J., Notin, P., Dias, M., Gomez, A., Min, J. K., Brock, K. P., Gal, Y., and Marks, D. S. Disease variant prediction with deep generative models of evolutionary data. Nature, 2021.
Grave, E., Cisse, M. M., and Joulin, A. Unbounded cache model for online language modeling with open vocabulary. In Advances in Neurul Information Processing Systems. vol- ume 30. Curran Associates, Inc., 2017. URL https: / //proceedinge wig.eiiinnsgcce www.eeeidinnnneeuiieeeeeeeeeeeee wwwp:eeuuurirr.crrreeeeeeeeeeeee wwtp:pppcpcppppppppccce hash/ (tchh//44ce2639990 Abstract html.
Guu, K. Lee. K. Tung. Z. Pasupat, P.. and Chang. M. Retrieval Augmented Language Model Pre-Training. In Proceedings of the 37th International Conference on Machine Learning, pp. 3929-3938. PMLR, November 2020. URL https: : /prpceeeinggg www.eeiiiinniiiin.ieeeeeeeeeeeee mir.press/ v119/guu200 htmml..
References
html. ISSN: 2640-3498.
Alley, E. C., Khimulya, G., Biswas, S., AlQuraishi, M., and Church, G. M. Unified rational protein engineering with sequence- deep representation learning. Nature Methods, pp. 1-8, 2019.
Göhel, U., Sander, C., Schneider, R., and Vaencia,, A. Corre- lated mutations and residue contacts in proteins. Proteins:
Structure, Function, and Hioinformatics, 18(4)3309-317, 1994, ISSN 1097-0134. doi: 10. 1w..102/poo.340080040000 URL doi/abs/10 1002 / prot wttpprttt...00 340180402. http:://niiiniiiiberyywiieeooooo
https:/ wttsss///nninneieaaaaee wwweyw wil.y.coml
Altschul, S. F., Gish, W., Miller, W., Myers, F.. W., and I. Lipman, D. J. Basic local aligmment search tool. Journal of Molecular Biology, 215(3):403-41 October 1990. ISSN 0022-2836. doi: 000066/800222886005)))80000
eeprint:
Haddox, H. K.. Dingens, A. S., Hilton, S. K.. Overbaugh. J. and Bloom, J. D. Mapping mutational effects along the evolutionary landscape of HIV envelope. eLife, 7:e34420, mar 2018 ISSN 2050-084X. doi: (..755//l.ifc.344220 URL httpt:tt/////t /doi. org/ 10. 7554/55111 fe. 34420.
Bridgford, J. L., Lee, S. M., Lee, C. M. M., Guglielmelli. P., Rumi, E., Pietra, D., Wilcox. S., Chhabra, Y., Rubin, A. F., Cazzola, M., Vannucchi, A. M., Brooks, A. J., Call, M. E., and Call, M.J. Novel drivers and modifiers of MPL--dpenndent oncogenic transformation identified by deep mutational scanning. Blood, 135(4):287-292. 01 2020. ISSN 0006-4971. blood.2019002561. URL https::/p//.....t.. doi.. org/10.1182/ blood. 022000202562
Hamamsy, T., Morton, J. T., Blackwell, R., Berenberg, D., Car- riero, N., Gligorijevic, V., Strauss, C. E. M., Leman, J. K., Cho, K., and Bonneau, R. Protein remote homology detection and structural alignment using deep learning. Nature Biotech- nology, 42(6))975-885, June 2024. ISSN 1546- 1696. doi: 10. 10.103884415577-023-009977 URL battps://doi. /org 10. 10011038//00 10111///41177-02-01111111 0211111---11112
doi: 10.1182/
Buchfink, B., Xie, C., and Huson, D. H. Fast and sensitive protein alignment using DIAMOND. Nature Methods, 12(1):59-66 60, 2015. doi: 10.18//mmehh.117
Hong, L., Sun, S., Zheng. L., Tan, Q., and Li, Y. fastMSA: Accelerating Multiple Sequence Alignment with Dense Retrieval on Protein Language, December 2021. URL https: / /www. wwwrxii www.:/etteeccteeeeeeeeee..eee.e. /oo.tt/1..11101. 202 12.20.. 473431v1. Pages: 2021.12.20.473431 Section: New Results
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. FlashAt- Fast and Memory-Ffficieet Exact Attention with TO- Awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
tention:
Douze, M., Guzhva, A., Deng, C., Johnson, J., Szilvasy G., Mazaré, P.EE., Lomeli, M., Hosscini, L., and Jégou, H. The Faiss library, September 2024. URL http:/ attt:://axxiv /arxiv.orgg/ abs/ 200s/401... 1. 08281. arXiv:2401 .08281 [cs].
Hopf, T. A., Schärfe, C. C. I., Rodrigues, J. G. L. M., Green, A. G., Kohlbacher, O., Sander, C., Bonvin, A. M. J. J., and Marks, D.S. Sequence co-evolution gives 3D contacts and struc- tures of protein complexes. el.ife, 3:e03430, September 2014. ISSN 2050-084X. doi: D..7554/cLiffe.3430 URL https:
Edgar, R. C. MUSCLE: multiple sequence alignment with high
10

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
//doi.org/// 70547554///e 03430. Publisher: eLife Sciences Publications, Ltd.
Karpukhin, V., Oguz, B., Min, S., Lewis, P., Wu, L., Edunov, S., Chen, D., and Yih, W.-t. Dense Passage Retrieval for Open-Domain Question Answering. In Webber, B., Cohn, T., He. Y., and Liu, Y. (eds.), Proceedings of the 2020 Confer- ence on Empirical Methods in Natural Language Processing (EMNLP), pp. 6769-6781, Online, November 2020. Association for Computational Linguistics. doi: 10. 10.18663/ main.550. URL https: http::://aaaaatin htt///aaaaaanntogogaaeeaa .rgg wwt.://oooogoog....o...o emnlp 1alp--main iamaiin.500m 550/.
Hopf, T. A., Ingraham, J. B., Poclwijk, F. J., Schärfe, C. P.I.I. I., Springer, M., Sander, C., and Marks, D. S. Mutation effects predicted from sequence co-variation. Nature Biotechnology, 35(2):128-135, February 2017. ISSN 1546-1696. doi: 10, 1038/ nbt.3769.
Izacard. G. and Grave, E. Leveraging Passage Retrieval with Gen- erative Models for Open Domain Question Answering, Febru- ary 2021, URL http: htpp:/aarxe http//awwcrggg//oo..////// abs 01282. arXiv:2007.0m 012282 [es].
Krogh, A. An introduction to hidden Markov models for bi- ological sequenees. In Salzberg, S. L., Searls, D. B., and Kasif, S. (eds.), New Comprehensive Biochemistry, volume 32 of Computational Methods in Molecular Biology, pp. 45- 63. January 1998. doi: 10.1016/50166--33666666666 htLpttps:/ htttpp::/www..ciennedieeeee www.weeeeeeccnecntec.ccco com/ science/ aet.cc/e/pi./s01177766666688
Izacard. G., Lewis, P., Lomeli, M., Hosseini, L., Petroni, F., Schick, T., Dwivedi-Yol Yu, J., Joulin, A., Riedel, S., and Grave, E. At- las: Few-shot Learning with Retrieval Augmented Language Models, November 2022. URL http:/ / //aaaa..vvv//////////// www.iiiorrraaaaao /bbs/ 2208. 03299, arXiv:2208. Si0.:22000002999999 [cs].
Lee, K.. Chang. M..W., and Toutanova, K. Latent Retrieval for Weakly Supervised Open Domain Question Answering. In Ko- rhonen, A., Traum, D., and Mårquez, L. (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp.
Jacquier, H., Birgy, A., Nagard, H. L.,, Mechulam, Y., Schmitt, E., Glodt, J., Bercot, B., Petit, E., Poulain, J., Barnaud, G., Gros, P.- A., and Tenaillon, O. Capturing the mutational landscape of the beta-lactamase TEM-1. Proceedings of the National Academy of Sciences, 130632) 13072, 2013. doi: 10.10as/pnas. 1215206110. URL bt bttps:/ www. pnas. org/ doi / aabs/ 10. 10..101///paa wtw.s///////.///////aa/////..a.. ppt.//paassaa....a 1215206110.
Florence, Italy, July 2019. Associa- tion for Computational Linguisties. doi: 10.18653/v1/P19 1612. URL https: / //a//aaaaaaat calantooooglyy org.
Lewis, P., Perez, E., Piktus, A., Petroni, F., Karpukhin, V., Goyal. N., Kütttle, H., Lewis. M.. Yih, W.-t, Rocktäschel, Gener- Advances
T.. Riedel, S., and Kiela. D. ation for Srwwledge-nntensing NLP Tasks. in Neural Information Processing Systems, volume 33, pp- Curran Associates, Inc., 2020. URL https: / /pproceedings www.insnuiipicoom co/ppe///cc//pppp//////// hash/ Abstract htt.l.
Johnson, J., Douze, M., and Jégou, H. Sillion-Scale Similarity Search with GPUs. IEEE Transactions on Big Data, 7(3); 535-547, July 2021. ISSN 2332-7790. doi: 10.1109/TDAAA. 2019.2921572. URL https: / //weeeepeelreeeeeeeeeee leee. orr// Document//873000 Conference Name: IEEE Transactions on Big Data.
Johnson, L. S., Eddy. S. R., and Portugaly, E. Hidden Markov model speed heuristic and iterative HMM search procedure. BMC Bioinformatics, 11(1))431, August 2010. ISSN 1471- 2105. doi: 10.1186/1441 17111100 100005-1-44- 11-431. URL https://dodo org/1 /00011111100 11 86/ 1 471 17771122 2 05 101-1111 411.
Li, P., Cheng. X., Song, I.., and Xing, F. Retrieval Augmented Protein Language Models for Protein Structure Prediction, December 2024. URL https:/ // /wwpw//ww.///////////// biorxiv.orop
cont / /10 1ee11111122220 .2. 02. 626519v1. Pages: 1.24.12.02.626519 Section: New Results.
Jumper. J.. Evans, R., Pritzel, A.. Green, T., Figurnov, M.. Ron- neberger, O., Tunyasuv K., Bates, R., Zidek, A., Potapenko, A., Bridgland, A., Meyer, C., Kohl, S. A. A., Bal- lard, A. Cowie, A., Romera-Paredes, B., Nikolov, S., Jain, R., Adler, J., Back, T., Petersen, S., Reiman, D., Clancy, E., Zielin- ski, M., Steinegger, M., Pacholska, M., Berghammer, T., Boden- stein, S., Silver, D., Vinyals, O., Semior, A. W., Kavukcuoglu, K., Kohli, P. and Hassabis, D. Highly accurate protein structure prediction with AlphaFold. Nature, July 2021.
Lin, Z. Akin, H., Rao, R., Hic, B., Zhu, Z. Lu, W., Smetanin, N., Verkuil, R., Kabeli, O., Shmueli, Y., dos Santos Costa, A., Fazel-Zarandi, M., Sercu, T., Candido, S., and Rives, A. Evolutionary prediction of atomic- level protein structure with language model. Science, 379(6637): 121-)111113113000 March 2023. doi: 10.1126/sceence ade2574. URL https:/ / //ww. science org/ doi/10. 1126/ stt//sciecc....c adeieeeeeeee233 2 74. 4. Publisher: American Associa- tion for the Advancement of Science.
Jégou, H. Douze, M., and Schmid, C. Product Quanti- zation for Nearest Neighbor Search. IEEE Transactions on Pattern Analysis and Machine Intelligence, 33(1):117-
Llinares-López F., Berthet, Q. Blondel, M.. Teboul, O., and Vert, J.PP. Deep embedding and alignment of pro- tein sequences. Nature Methods, 20(1):104-11,, January 2023. ISSN 1548-7105. doi: 10110386411922-02--011 2. URL https:/ / //www www.:natur..ecmm com/ arrtiritii ictes e21592-022-01700-12 Publisher: Nature Publishing Group.
128, January 2011. ISSN 1939-3539. doi: 10.1 TPAM1.2010.57. URL https: / org/ htss//occmett... /4/2202222222 Conference Name: IEEE Trans- actions on Pattern Analysis and Machine Intelligence.
/ttpe////eeeeell///o/oeeeeoo eeexplore leee.
Kallenborn, F., Chacon, A., Hundt, C. Sirelkhatim, H., Didi, K., Cha, S., Dallago, C., Mirdita, M., Schmidt, B., and Steinegger, M. GPU-accelerated homology search with MM- seqs2, January 2025. URL https: //www. www.bii.oii...............i.iiwi wwxxviorrrg org/ content/ / /10. 21101202411. 11_111111 13. 623350v3. Pages:
Ma, C., Zhao, H., Zheng, L., Xin, J., Li, Q.. Wu, L., Deng, Z.. Lu, Y. Y., Liu, Q. Wang. S., and Kong. L. Retrieved Se- quence Augmentation for Protein Representation Learning, In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. (eds.), Proceed- ings of the 2024 Conference Empirical Methods in Nat- ural Language Processing, pp. 1738-1767, Miami, Florida,
2024.1 12.63330 Section: New Results.

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
USA, November 2024, Association for Computational Linguis- tics. doi: 10. 10.18633/1220244cmmmmmmmmm 104. URL https: //aclanthologe: wrg.
with differentiable Smith-Wait Waterman. Bioinformatics, 39 (1):btac724, January 2023. ISSN 1367-4811. doi: 10.
1093hhoonformaiis//aac77444 https:t::/ htttt:::/ddc..org///// http://looonoormaaicssaaaaaaaaed
Madani, A.. McCann, B., Naik, N., Keskar, N. S., Anand, N., Eguchi, R. R., Huang, P.-.,, and Socher, R. ProGen: Language Modeling for Protein Generation, 2020. _eprint: 2004.03497.
Potter, S. C., Luciani, A., Eddy, S. R., Park, Y., Lopez, R., and Finn, R. D. HMMER web server: 2018 update. Nucleic Acids Re- search, 46(W1): 1):W200-W222 July 2018. ISSN 0305-1048, doi: 10ar URL httph::t/////e http://///o...o/.....////// nar/ /gy4/88
McLaughlin Jr, R. N., Poelwijk, F.J. J., Raman, A., Gosal, W. S., and Ranganathan, R. The spatial architecture of protein function and adaptation. Nature, (914722): 2012. doi: 10. 1038/nature11500
Rao, R. M., Liu, J., Verkuil, R., Meier, J., Canny, J., Abbeel, P., Sercu, T., and Rives, A. MSA Trans- former. In Proceedings of the 38th International Con- ference on Machine Learning, pp. 8844 8856. PMLR. July 2021. URL https: / pttpppeeeeeggsss...oe// 2139/rao2la.hie
Mcier, J., Rao, R., Verkuil, R., Liu, J., Screu, T., and Rives, A. Language models enable zero-shot pre- diction of the effects of mutations on protein func- tion. bioRxiv, 2021. doi: 00.100/2021.00.00..5500000 htth::// /www. biorxiv. .owg/ htt////onnee/trtryy 2021/07/10 011//112222..1 07. 09, 450648. Publisher: Cold Spring Harbor Laboratory.
www.iinnss...rrreeeeeeeee mir.press/ ISSN: 2640-3498.
Riesselman, A. J., Ingraham, B., and Marks, D. Deep genera- tive models of genetic variation capture the effects of mutations. Nature Methods. 15(10):816 1110))816-822, 2018. Publisher: Nature Publishing Group.
Nambiar, A., Heflin, M., Liu, S., Maslov, S., Hopkins, M., and Ritz, A. Transforming the Language of Life: Transformer Neural Networks for Protein Prediction Tasks. In Proceed- ings of the 11th ACM International Conference Bioinfor- matics, Computational Biology and Health Informatics, BCB '20, pp. 1-8, New York, NY, USA, November 2020. Associa- tion for Computing Machinery. ISBN 978-1-4503--964-99 doi: 10. (0.1145533884400344244444 URL https::tt...... : /dl. acm. doi/10.1145 0384444//33344444 3311//33388400.31145777 3412467.
Riley. A. C., Ashlock, D. A.. and Graether. S. P. The difficulty of aligning intrinsically disordered pro- tein sequences as assessed by conservation and phy- logeny. PLOS ONE, 18(7)xc0283388 July 2023. ISSN 1932-6203. doi: 10. http://rnalpoo..0228s8 URL https: /atee///jouraaaa s. plos. org/plosone/ article?id=10- 112100113711 1371/ / pone 0288388. Publisher: Public Library Science.
Nijkamp. E., Ruffolo, J. A.. Weinstein, E. N., Naik, N., and Madani, A. ProGen2: Exploring the boundaries of protein language models. Cell Systems, 14((1))988-97) e7, November 2023. ISSN 2405-4712, 2405-4720
Rives, A., Meier, J., Sercu, T., Goyal. S., Lin, Z., Liu, Guo, D., Out, M., Zitnick, C. L., Ma, J., and others. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences, 1 118(15), 2021. Publisher: National Acad Sciences.
doi: 10. 1016/f 10. 002. URL https:://///....... / www. ccll. / abbsraatts224555 (23) 00272-20000222222 vier.
Publisher: Else-
Notin, P., Dias, M., Frazer, J., Marchena-Hurtadoe J., Gomez, A., Marks, S., and Gal, Y. Tranception: protein fitness pre- diction with autoregressive transformers and inference-time time re- trieval, May 2022. URL ht / htt///axvvvvoro///ee/ee / abs 2205. 13760. arXiv:2205.1 13760 [es].
Robertson. and Zaragoza, H. The Probabilistic Relevance Frame- work: BM25 and Beyond. Found, Trends Inf. Retr. 3(4):333- 389, April 2009. ISSN 1554-0669. doi: 10.1561/15000000 URI htw.hhtt////// //doi org/ /0////1.111 1 1561 500000019 19.
Notin, P., Kollasch, A., Ritter, D., van Nickerk, L., Paul, S., Spinner, H., Rollins, N., Shaw, A., Orenbuch, R., Weitzman, R., Frazer, J.. Dias, M., Franceschi, D., Gal, Y., and Marks, D. ProteinGym: Lcrge-Scale Benchmarks for Protein Fitness Pre- diction and Design. Advances in Newral Information Prucessing Systems, 36:64331 64379, December 2023. URL https: aipapers..nig cc http://ccccccpceeees///e//////// hach:ccac7233511122666333338 htt.rracc--aatasetss_s.
Roscoe, B. P.. Thayer, K. M., Zeldovich, K. B., Fushman, D., and Bolon, D. N. Analyses of the Effects of All Uhiquitin Point Mutants on Yeast Growth Rate. Jour- nal of Molecular Biology, 1237 2013. ISSN 0022-2836. doi: hutp:///oioag///0.00....mm...... URL https: htt.ttts///ww.ccieeeeee http://wwwwcieeeeeeeeeeeeee com/ wttm///ccccc///////cc//c/ htt.cc/ee/pps00222222111 wtt.a///ppp//0022281130000000000
html.
Russ, W. P., Figliuzzi, M.. Stocker, C., Barrat-Charlaix, P., So- colich, M., Kast, P., Hilvert, D., Monasson, R., Cocco, S., Weigt, M., and Ranganathan, R. An evolution-based model for designing chorismate mutase enzymes. Science, 369:440 445, 2020. URL ht tps::ps//s / //ap.. swwant icschola.. org/CorpusID 220714458.
Notin, P., Rollins, N. J., Gal, Y., Sander, and Marks, D. Machine learning for functional protein design. Nature Biotechnology, 42: 216-228, 2024, URL https:tt.. / /t///apa/ /api. www.steeeiiiii icschoaa..ccc..r..... org/ Org/Coo http/Corpusiic ID: 27276627799
Olson, C. A., Wu, N. C., and Sun, R. A Comprehensive Bio- physical Description of Pairwise Epistasis throughout an Entire Protein Domain. Current Biology, 24(22):664-2655 2014.
Sachan, D.,, Patwary, M., Shocybi, M., Kant, N., Ping, W., Hamil- ton, W. L., and Catanzaro, B. End-to-Ena Training of Neural Retrievers for Open-Domain Question Answering. In Zong, C., Xia, F., Li, W., and Navigli, R. (eds.), Proceedings of the 59th Annual Meeting of the Association for Computational
Petti, S., Bhattacharya, N., Rao, R., Dauparas, J., Thomas, N., Zhou, J., Rush, A. M., Koo, P. and Ovchinnikov,
S. End-to-end learning of multiple sequence alignments
12

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
Linguistics and the 11th International Joint Conference on Nat- ural Language Processing (Volume I: Long Papers), pp. 6648- 6662, Online, August 2021. Association for Computational Lin- guistics. 10.186533/ acc Ion---ong,,19, URL https: h/walanthooogyy....... www.//oy....o............. wrg. 2021 aal..llccll -ong. 519//
Urbina,
F. Lentzos, F., Invernizzi, C., and Ekins, S. suiffraal-melllienne--ooeer drug discov- ery. Nature Machine Intelligence, 4(3):189-191, March
Dual
use
2022.
ISSN 2522-5839. doi: 10. 1000864225-022-0006600
9. URL https: /www nature. com/ artiileti
a42256- 022-0065--9 Publisher: Nature Publishing Group.
Sgarbossa, D., Malbranke, C., and Bitbol, A.-F ProtMamba: homology-ooo aware but alignment-free protein state space model, May 2024. URL https: /wwww//www...../w/////////.. www.boiioivvoo...........ooo.. content 10. 110112024, wwt..................... 595730v2. Pages: 2025..5.24..55533 Section: New Results.
Wang, S., Yu. M., Guo, X., Wang, Z., Klinger. T., Zhang, W. Chang, S., Tesauro, G., Zhou, B., and Jiang. J. R3: re- inforced ranker-reader for open-domain question answering. In Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence and Thirtieth Innovative Applications of Artificial Intelligence Conference and Eighth AAAI Sym- posium on Educational Advances in Artificial Intelligence, AAAI' DAI'18/1A Louisiana, USA, February 2018. AAAT Press. ISBN 978-1- 57735-800-8.
Shu, Y. and McCauley, J. GISAID: Global initiative sharing all influenza data from vision reality. Euro- surveillance, 22(13)330494 March 2017. ISSN 1560-7917. URL https: http///wwrrrsureelannee...ceeeee ootpssarreeeiaaeeeeooooooottttt. org/content/
doi: 10.287/1150-77171852202222
pp. 5981-5988, New Orleans,
00.22807116-0 7917.E3 2017.22 22. 1..22.1..0... 30494.
Publisher: European Centre for Discase Prevention and Control.
Weile, J.. Sun, S., Cote, A. G., Knapp, J., Verby, M., Mel- lor, C., Wu, Y.. Pons, C., Wong, C., van Lieshout, N., Yang. F, Tasan, M., Tan, G., Yang. S., Fowler, D. M., Nussbaum, R., Bloom, J. D., Vidal, M.. Hill, D. E., Aloy, P. and Roth, F. P. A framework for exhaustively map- ping functional missense variants. Molecular Systems Biol- http://doi..r/100...5233 egy, 13(12):957, 2017. doi: embooresss org/ doi 20177908. URI. https: : //wwwww.. http://ww..op...eeeeeeeeeee abs/10. 1 15552/mmbb 20177908.
Starr, T. N., Greaney, A. J., Hilton, S. K., Ellis, D., Crawford,
K. H., Dingens, S., Navarro, J., Bowen, J. E., Tortorici, M.A., A., Walls, C., King, P., Veesler, D., and Bloom, D. Deep mutational scanning of SARS-CoW receptor binding domain reveals constraints on folding and ACE2 binding. Cell, 1295-11199 2020.
Steinegger, M. and Söding, MMseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature Biotechnology. 35(11) 1028-1028, November 2017. ISSN 1546- 1696. doi: 10. 1088. URL https:://// /www. nature. com/articlese nbb. 3988. Publisher: Nature Publishing Group.
Wu, N. C., Olson, C. A., Du, Y., Le, S., Tran, K., Remenyi, R., Gong, D., Al-Mawsawi, L. Q. Qi, H., Wu, T.-T,, and Sun, R. Functional Constraint Profiling of Viral Pro- tein Reveals Discordance of Evolutionary Conservation and Functionalityy. PLOS Genetics. 11(7):1-27, 07 2015. doi: 10. 1/jurnajj 1005310. URL httpsstts //o..orgg 10. 1371 //111//ouu// journall pgen. 1005310.
Suzek, B. E., Wang. Y., Huang, H., McGarvey, P. B., and Wu, C. H. UniRef clusters: comprehensive and scalable alter- native improving sequence similarity searches. Bioinfor- matics, 31(6):26--32,, March 2015. ISSN 1367-4803. doi: 10. 1...0///iiiinfmmatiiss/u33399999 URL https: : /wwww ncbi. nim. nih. http:///cc//atticcee///////033//
Thompson, J.
D., Higgins, D. G., and Gibson, T. CLUSTAL W: improving the sensitivity of progressive multiple sequence align- ment through sequenee weighting, sppecfi gap penal- ties and weight matrix choice. Nucleic Acids Research, 22(22); 4673-4680, 11 1994. ISSN 0305-1048, doi: 10. 10.1099/mm/22..2.. 4673. URL https:/ /ooi.oioo/////o/ org/10 10..111..11..0.0 1093 4673.
Tripathi, A., Gupta, K., Khare, S., Jain, C., Patel, S., Kumar, P., Pulianmackal, A. J., Aghera, N., and Varadarajan, R. Molecular Determinants of Mutant Phenotypes. Inferred from Saturation Mutagenesis Data. Molecular Biology and Evolution, 33(11); 2960-2975, 08 2016. ISSN 0737-4038. doi: 10. 1093/mllhew/ 1sw182. URL https: http:///ooooo//////// 1...//1..11000 msw182.
1093/ 1093//0mme/////////moooo /bevv
Truong Jr. T. and Bepler, T. PoET: A generative model of protein families as soguenccs-of-sequeeee
Ad-
vances in Neural Information Processing Systems, 36:77379-77415, December 2023. URL https: //papers wttpppssccppppeeffeeee//eeeeeeee
2et.34666eeeee22222eeeeeeeeee 93ccabb Abstract- Celtraat-Cooffee Conference.cccee html.
13

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
Appendix
A. Vector similarity search with Faiss
We rely on Faiss for GPU-accelerated vector similarity search (Johnson et al., 2021; Douze et al., 2024), whose terminology we adopt. The vector similarity search is facilitated with index, whose task it is search large database of vectors, d and return the K most similar ones the query, q. given similarity metric.
The most simple index is a flat index, where the query compared to all database entries. With the commonly used maximum inner product similarity measure, this reduces to computing q d" and extracting the K largest entries. While this search is exact, it is both slow and requires storing all database vectors in memory which is prohibitively expensive. For fast search, an inverted file index (IVF) can be used. Prior searching the index, all entries are clustered using "coarse quantizer", e.g.. k-means clustering algorithm, given some predefined number of centroids, Kiv-- At search time, the query q is compared all Kivp centroids, of which the Piyr most similar centroids, often referred to as the number of probes, are searched, reducing the number of comparisons from N to
N KIvF
Noomparisons KIVF + Plvf
as per equation 18 in Douze et al. (2024). While using an IVF index reduces search time, it still requires storing the full database in memory, which for the as 62 million UniRef50 sequences requires > 110 GB of memory, given the 480 dimensional mean-pooled ESM-2 35M embeddings. To overcome this major challenge, further quantization is required. We rely on a product quamtizer (PQ) effectively reduce the dimensionality of each vector (Jégou et al., 2011). The product quantizer partitions each vector into M sub-vectors, where each sub-vector is further separately quantized using a k- k-means clustering. Defining the product quantizer requires setting two parameters: the code size, M. and the number of bits with which to represent each sectvec,or, where either 8 or 10 are commonly used.
Using product quantizer dramatically reduces the index size. The memory requirement of indexing UniRef50 using three different code sizes and using a flat IVF index be seen in Figure A. IVFPQ32x8 refers product-quanized IVF index, where each vector divided into M 32 sub- aub-vretors, each of which is represented by bits. The shown memory uses are solely for storing the index memory. Using a quantizer such as PQ is therefore necessary in order to additionally store and train the Protriever model.
The process of preparing the coarse and product quantizers, c.g., by running k-means algorithms to facilitate fast search is called training the index.
133.9
125
Memoory (GiB)
75
50
25
8.9
9.8
12.6
IVEPQ32x8
IVHPQ48X8
IVEPQ96x8
IVFFlat
Figure I. Index memory use. TVFPQ32x8 refers product-quantized IVF index, where each vector is divided into M 32 suh- vectors, each of which is represented by 8 bits, while IVFF lat refers to an IVF index with no further quantization. Using a product quantizer dramatically reduces the memory use, potentially at the cost of search quality.
A.1. Distributed index
We use Protriever in a distributed setting using Naru GPUs via the implementation in Izacard et al. (2022). We shard our datasct into Neru equal-sized partitions and train separate indices, where each index is responsible for Numex N/NoPu sequenees. At search time, the query is used to search each index, the results of which are aggregated. We let Naru
A.2. Choosing index parameters
We need to select a number of index parameters, namely the number of centroids for the coarse quantizer, Kivr, the number of probes, Pive, the code size M. and the number of bits for the product quantizer. We have three main considerations: memory, speed, and accuracy. Given the memory uses shown Figure 1, now focus on gauging accuracy and speed.
A. RECALL
We measure search accuracy using recall by randomly sampling Nample 10000 UniRef50 sequences queries and investigating whether the query sequences are returned when searching the index. We investigate code sizes of 32, 48, and 96 (the embedding dimension
14

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
needs divisible by the code size), coarser quantization led poor performance. We experiment with three different centroid
counts, determined by database size: Kivy E { [VNindex: NyNndsss 8yNmite). We fix the number of probes to Pivr = 2048, which the upper limit in the Faiss GPU implementation and for simplicity, we fix the number of bits per sub-vector to 8. This lcads to three indices IVFFQ32x8. IVFPQ4 8x and IVEPO Px96,,, with KIVE [3941, 15764, 31528] (as Ninexx 15.5 million). We use the 10,000 sampled queries to search across the nine index configurations, retrieving the K 2048 nearest neighbors and calculating the average recall rate at powers
results can be seen in Figure A.2.
0.6
Recall
0.4
Index
[VE803228
TVFPQ45<8
[VTPQ9668
0.2
Kxy
3941
15764
31528
16 32
64
128
256
512
1024
2048
Figure A.2. Recall rate vs. neighborhood sizes for IVEPQ indices at different quantization levels centroids counts. 10,000 UniRef50 sequences are randomly sampled and used queries. For each query sequence, the 2048 nearest neighbors are found. The recall indicates whether the query sequence was successfully recovered. Decreasing the quantization from 48 sub-vectors 96 sub-vectors leads to significant increase in recall, while doubling the number of centroids per index from Kivr 15764 to KivE 31528 only has a marginal performance increase.
The code size has significant impact on search quality, where = 332 (shown bluc) fails to reach recall rates above 0.9 for K 2048, Increasing the code size to M = 96 significantly increases the recall rates, where the majority of the single-ncarest neighbors return the query sequence. The search performance is less sensitive to the number of centroids, where recall increases with centroid count. For M 96, observe persistent performance gap when using KivF / 15764, particularly for the lower neighbor counts.
A.2.2. SEARCH SPEED
We investigate impact parameter settings on search speed by measuring it over a range scenarios. each the nine parameter configurations, we search using 1, 10, and 100 queries, repeating each five times. We can visualize the search time (averaged over repeats), the time per query, and the queries expected, the search time increases with the number of queries. Using low numher of centroids (shown in blue) consistently leads longer search times. While the search process has fewer comparisons to centroids initially, this is outweighed larger clusters. the second row Figure A.3 we observed the search time per query longer when using only single query. This expected as Faiss is optimized for batched searches. We also observe that using code sizes 48 and 96 approximately takes the same scarch time per input query. In the last row we observe that the queries second (QPS) generally increases with the number of queries and for code size 96 appears near saturation point.
second (QPS) for each configuration. These results can he seen in Figure A.3. We see that,
the correspondingly
A.2.3. INDEX TRAINING TIME
We lastly examine how Jong it takes to train the index with the different parameter configurations. We train each of the nine configurations on Ninde 15.5 million UniRef50 sequences a total three times. The average training time in seconds standard error can be seen Tahle A.1. The training time is not sensitive to the code size hut appears to linearly scale with number centroids.
A.3. Retrieval time comparison with MMseqs2
We here provide the benchmarking procedure used to generate the results in Figure 2. The tabulated results (mean and standard deviation) additionally shown in Table A.2. We define query sizes of 1, 10, 100, and 1000 and sample total of five query from UniRef50 random for each size, resulting in 20 distinct query sets. These are then used as queries for Protriever (i.c., vector similarity search
15

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
IVFPQ32X8
IVEPQ48X8
IVFPQ96X8
12.5
10.0
KivF
Time (8)
75
3941
15764
31528
2.5
0.05
Time (uery
0.04
KiVF
3941
15764
31528
0.01
0.00
Queries per second (QPS)
300
2.0
Kivr
200
3941
1.90
15764
100
31528
Number queries
Number queries
Number of querios
Figure A.3. Impact of parameter settings search time metrics (total time, time per query, and queries second) across 9 configurations with varying number queries. Results show longer search times with fewer centroids, higher per-query costs for single queries due to batch optimization, and QPS improvements that scale with number of queries..
16

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
Training time (s)
Kivp = 3941
KIvF 15764
Kive = 31528
M 32
39.00 + 1.36
80.63 + 0.47
190.13 0.38
M 48
40.33 + 0.52
82.89 + 0.73
191.72 1.59
M 96
47.03 = 0.60
89.98 = 0.84
200.13
1.69
Table A. Index training times. Average index training times (and standard error) for different parameter configurations. Each index covers Nindex 15.5 million UniRef30 sequences. The indexing time only slightly decreases with increased quantization. The number of centroids has a large impact on indexing time which appears to scale linearly.
through Faiss) and MMseqs2 in both CPU and GPU modes. Each search is repeated five times to account for runtime variability. To process the results, we take the median over repeated runs and compute the mean and standard deviation the distinct query sets.
We use the following search parameters for MMseqs as described Kallenborn et al. (2025) and detailed the user guide:
mmseqs search queryDB targetDB resultDB tmp -s 8.5 -e 10000 --max- -ses 4000
--prefilter-
1 --threads 20
for the CPU-version and
mmseqs search queryDB targetDB.qpu resultDB tmp --epu --gpu-server --db-load-mont 2 --threads 10000 --max-segy 4000
for the GeU-vorsion.
Both query and database FASTAs are converted into MMseqs2 databases prior to searching. We additionally pre-compute the target database indices. For MMseqs2-GPU, we follow the guidelines and spawn a gpuserver via mnseqs gpuserver target gpp --max- seqs 4000 prior to searching for reduced overhead. For Protriever, the query sequences are embedded prior to searching. The reported results thus show the time it takes for the search itself, assuming that appropriate pre- processing has been conducted for all three methods.
Protriever GPU-accelerated MMseqs2 searches are made on single L.40S GPU using one CPU thread. The CPU-version MMseqs2 uses 20 threads due to compute constraints. The Protriever Faiss index is pre-trained on all UniRef50 sequences and does not use sharding for direct single-GPU comparison.
Retrieval time per query (s)
N
N 10
N 100
N = 1000
Protriever (Faiss)
0.007 = 0.0005
0.005 0.0001
0.005 + 0.0001
0.005 = 0.0001
MMseqs2-GPU
0.637 + 0.3079
0.532 1.12
0.493 + 0.0653
0.508 0.0246
MMseqs2 (k-mer)
28.587 + 5.7237
17.771 2.9081
15.010 + 1.4492
15.567 + 1.0836
Table A.2. Retrieval time per query. Average retrieval time (and standard deviation) per query function of number of queries over various query sets.
A.4. Retrieval homology
To investigate the homology of the retrieved sequences, we compare Protriever to MMseqs2 (Steinegger & Söding, 2017; Kallenborn al., 2025) and JacKHMMER (Potter et al., 2018). We use the ten validation sequences in ProteinGym (Notin et al., 2023) queries for the retrieval methods. For Protriever and MMsequ2, use the settings detailed in Appendix A.3, where for JacRHMMER, we use scores of 0.1, 0.5, and 0.9. We compute the sequence identity between query sequences and all retrieved sequences their distribution in Figure A.4. As JackHMMER returns a highly variable number of sequences, e.g., a single UniRef50 sequence was returned using IKIE SARS2 2 query with bit score 0.9, while 9.839 sequences were returned for CALM1 HUMAN with bit score 0.1, we weigh the distribution by the number of retrieved sequences for each method/query pair. This way each of the 10 retrieval sets contribute equally to distribution visualization. During inference, retrieved sequences with than % sequence identity are removed. The resulting distributions post-filtering are shown in Figure A.5, where we now observe gap between MMseqs2 and Protriever, where the latter retrieves sequences that are more similar to the queries.
visualize
17

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
0.08
Protriever (Faiss)
JackHMMER (bit score 0.1)
0.07
MMseqs2-GPU
JackHMMER (bit = 0.5)
MMseqs2 (k-mer)
JackHMMER (bit score 0.9)
0.05
Density
0.02
0.00
20
40
60
80
100
Sequence identity between queries and retrieved sequences
Figure A.4. Distribution of sequence identities. Distribution of sequence identities between queries and retrieved sequences for Protriever, MMSeqq2, and JacKHMMER, using the ten ProteinGym validation sets as queries. The retrieved sequences have not been filtered and are weighed by number retrieved sequences per method/query pair.
0.14
Protriever (Faiss)
JackHMMER (bit score
MMseqs2-GPU
JackHMMER (bit (bore - 0.5)
MMseqs2 (k-mer)
JackIIMMER (bit score 0.9)
0.10
Density
0.04
0.00
40
60
100
Sequence identity belween queries and retrieved sequences
Figure A.5. Distribution of sequence identities after filtering. Distribution sequence identities between queries and retrieved sequences for Protriever, MMSeqs2, and JackHMMER, using the ten ProteinGym validation sets queries. The retrieved sequences have been filtered by removing those with < 15% sequence identity to the query. The retrieved sequences are weighed by number retrieved sequenees per method/query pair.
18

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
B. Alternative architecture for the reader
This section presents results using Dusion-in-Decodee (FiD) alternative reader architecture, demonstrating the versatility of our retrieval framework across different architectural choices.
B.1. Fusion-in-Decote architecture
The Fusion-in- -nte--dee- (FiD) model (Izacard & Grave, 2021) integrates multiple retrieved sequences by independently encoding each one and fusing their representations during decoding. Unlike PoET's approach of processing concatenated sequences. FiD encodes each retrieved sequence separately through encoder, then concatenates their hidden representations for the decoder attend over.
Formally, given N retrieved sequences per query, each sequence is encoded independently to produce hidden that concatenated along the sequence dimension. During decnding, cross-attention allows the decoder to attend over all retrieved sequences simultaneously, effectively fusing information from the entire retrieved set.
Regarding computational complexity, FiD offers advantages over approaches that process all sequences jointly While PoET processes all sequenees simultaneously with quadratic complexity ((((k+1))) for l sequences of length 1, FiD's complexity scales O + +1% where sequences are processed independently before cross-attention computation. This linear scaling with the number of sequences makes FiD more computationally efficient, though potentially less expressive than architectures that model interactions between all positions across all sequences simultaneously.
We implement FiD using an ESM encoder (35M parameters) and Tranception decoder (85M parameters), connected through cross- attention layers, resulting in 150M parameter model. This approach is agnostic to the particular choice of encoder and decoder architectures, requiring only dimension alignment between encoder and decoder through projection layers.
pretrain this FiD model on same dataset used PoET training, following identical sampling and data augmentation strategies. We train our model (without DPR pretraining) within the Proriever framework, with EMDR end to end loss on the retriever for 50,000 iterations. We use A with batch size of 16, context size set to 20, learning rates of 4 x 10 - for the reader and 5 x 10 for the retriever, with lincar decay and 1,000 warm-up steps. We re-index dataset every 5, 5,000 steps for total of 10 re-indexing stages.
B.2. Fitness prediction performance of Iusion-in-Decoder architecture
We evaluate first the model with same MMseqs2 MSAs generated from uniref used in PoET, and using the same sampling scheme, based on (Hopf ct al., 2017XFiD + MSA in Table B.1) We also evaluate this trained MSA with Frozen ESM retriever, which substantially degrades performance, model with the trained Protriever retriever (0.416). We in particular better than the pretrained model evaluated with MSAs low depth. The results Table B. demonstrate of comparable size (FiD with trained Protriever at 150M parameters achieves 0,416 vs ESM2-M's 0.388) and achieve performance comparable to models that orders of magnitude larger (matching the 3B parameter ESM2-XL's 0.418 overall Spearman correlation), representing 20-fold reduction in model size while maintaining competitive fitness prediction accuracy. Note that these results preliminary. We could further improve these using DPR pretraining. which showed synergistic effects.
average going from 0.42 to 0.342. However, are able recapitulate the performance of the MSA input
retrieval-augmented models can significantly outperform single-sequence models
Table B.I. Zero-shot substitution DMS benchmark by MSA depth. Average Spearman's 's rank correlation between model scores and experimental measurements by MSA depth on the ProteinGym substitution benchmark. Tranception models are without inference-ting MSA retrieval. Alignment depth is detined by the ratio of the effective number of sequences Nor in the MSA. following (Hopf al., 2017), 7), by the length covered L (Low: No / L. <1; Medium: 1< Non / < 100; High: No# Tet/T. L. >> 100). The All column is the average across the three depths.
Model type
Model name
# Params
Spearman by MSA depth High
Low
Medium
All
ESM2-S
35M
0.239
0.271
0.453
0.321
ESM2-M
150M
0.306
0,358
0.500
0.388
Encoders
ESM2-L
650M
0.335
0.406
0.517
0,419
FSM22XIL
3B
0.348
0.415
0.491
0.418
Tranception-S
85M
0,258
0.295
0.321
0.291
Decoders
Tranception-M
300M
0.293
0.349
0.382
0.341
Tranception-L
700M
0.358
0.371
0.417
0.382
FiD + MSA
150M
0.352
0,411
0.498
0.420
FiD
FiD + frozen Protriever
150M
0.287
0.354
0.386
0.342
FiD + trained Protriever
150M
0.365
0.401
0.483
0,416
19

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
C. Detailed fitness performance results
C.1. Protriever results by MSA depth
Table C.I. Zero-shot performance segmented by MSA depth on the 217 substitution DMS of ProteinGym. Alignment depth
defined by the ratio of the effective number of sequences Ner in the MSA, following (Hopf al., 2017), by the length covered L (Low: Naw/ L <1; Medium: < Not/ <100; High: Nor/ > 100). designates Spearman rank correlation
Low MSA depth
Medium MSA depth NDCG
High MSA depth
Model type
Model name
NDCG
NDCG
Site independent
0.427
0.747
0.376
0.747
0.317
0.770
Alignment- based
GEMME
0.446
0.761
0,474
0.778
0.493
0.809
EVE
0.420
0.757
0.457
0.783
0.477
0.821
ESM-1v
0.316
0.685
0,409
0.743
0.495
0.808
Unconditional
ProGen2
0.323
0.727
0.412
0.775
0.442
0.808
PLM
ESM2
0.336
0.703
0.423
0.759
0.485
0.808
MSA Transformer
0.375
0.754
0,456
0.776
0.479
0.815
Tranception L
0.421
0.762
0.443
0.778
0.471
0.812
Hybrid
TranceptEVE
0.436
0.764
0.472
0.785
0.490
0.824
PoET
0.478
0.766
0,478
0.781
0.510
0.827
Protriever
Protriever
0.464
0.772
0.498
0.781
0.512
0.831
C.2. Protriever results by Taxa
Table C.2. Zero-shot performance segmented by Taxa the 217 substitution DMS of ProteinGym benchmark. spearman rank correlation.
designates
Human
Other Eukaryote NDCG
Prokaryote
Virus
Model type
Model name
P
NDCG
P
NDCG
NDCG
Site independent
0.380
0.759
0.389
0.781
0.318
0.770
0.375
0.695
Alignment-
GEMME
0.469
0.779
0.516
0.805
0.467
0.816
0.472
0.743
based
EVE
0.454
0.784
0.495
0.810
0,457
0.827
0.434
0.742
ESM-Iv
0.458
0.770
0.464
0.768
0.413
0.797
0.294
0.641
Unconditional
ProGen2
0.386
0.772
0.458
0.791
0.418
0.822
0.402
0.718
pLM
ESM2
0.442
0.778
0.477
0.775
0.458
0.814
0.294
0.652
MSA Transformer
0.439
0.780
0.516
0.812
0.446
0.823
0.421
0.723
Tranception L
0.455
0.788
0.497
0.807
0.414
0.812
0.438
0.727
Hybrid
TranceptEVEL
0.473
0.787
0.513
0.816
0.455
0.831
0.461
0.743
PoET
0.482
0.781
0.541
0.827
0.464
0.829
0.491
0.744
Protriever
Protriever
0.480
0.788
0.542
0.811
0.492
0.845
0.516
0.744
20

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
D. Protriever training loss function
In addition to the EMDR loss used in our main experiments (Equation (3)), we evaluated two alternative loss functions for end-to-end retriever training.
Perplexity Distillation (PDist) The PDist approach (Izacard et al., 2022) trains the retriever to predict how much each sequence improves the language model's perplexity when reconstructing the query sequence. We minimize the KL-divergence between the retriever's relevance scores (Equation (2)) and the posterior distribution based on language model performance:
exp Plog (q dx))
pk Σ exp (Log PLM (q d.))
(4)
Leave-One-Out Perplexity Distillation (LOOP) LOOP measures each sequence's contribution by evaluating how much the language model's reconstruction performance degrades when removing individual sequences from the retrieved set. For each retrieved sequence de, the relevance score defined as:
PLoOF (dx qq) = exp log (q Dx {dx)))
(5)
exp (- FLM PLM (q Dx \ {d.)))
While computationally more expensive than PDist and EMDR to requiring separate forward passes, LOOP better reflects the multi-sequence conditioning used during training by evaluating the model on (K -11 sequences rather than single sequences. Indeed, we observe in D.1 that LOOP performs the best with our FiD model in the Protriever framework. We however decided to pursue the use of EMDR for larger scale results with PoET due to its lower computational cost while maintaining competitive performance, Further work could look into LOOP alternative EMDR for the PoET Protriever model.
Table D. Spearman on validation set (Appendix F) for different losses with the FiD model. evaluate the FiD model with retrieved sets, sampled with the same scheme described in the main text. EMDR performs slightly better than the PDist loss. LOOP performs slightly better than the other two, requires many more forward passes
Training strategies
EMDR
PDist
LOOP
Frozen ESM
0.347
0.347
0.347
Protriever w/o DPR
0,404
0.397
0.409
E. Inference-time sampling of retrieved sequences
Our inference pracedure involves several steps ensure diverse and representative sequence sampling. We encode the query sequence and search the UniRef50 index for homologous sequences. From returned set, we filter sequences that have less than 15% sequence similarity with the query. We then sample 2560 UniRef100 sequences from the filtered UniRe150 homologous set, with weights inversely proportional to the size of their corresponding UniRet90 clusters. These sequences are encoded by the retriever and subsequently clustered using k-meaas with & 50.
Clusters are then sampled using weights /3- (1+e=)//T) where s is the size of the k-mean cluster and d is the distance to the query. Different combinations of and T values give rise to five parameter combinations that represent different trade-offs between diversity and relevance in the retrieved set. For experiments in Section 5.2, clusters are taken in UniRef50, and sampled with
same weights v (1+e- e-d/T)) -1 where d is taken as the alignment identity of the UniRe/50 cluster representative to the query, and is the size of the UniRet50 cluster. We also vary the conditioning set size across three configurations (6,144, 12,288, and 24,576 tokens). Overall, we ensemble across total configurations: sequence diversity strategies x 3 conditioning set lengths, with each configuration requiring separate forward pass through the model.
F. Validation sets from ProteinGym
We use the same ProteinGym assays for validation as in Tranception (Notin et al., 2022) and PoET (Truong Jr Bepler, 2023):
LEOOX (Jacquier et al., 2013)
DLG4_RAT (MeLaughlin Jr et al., 2012)
CALM1 HUMAN (Weile et al., 2017)
PA 134A1 (Wu et al., 2015)
CCDB ECOLI (Tripathi et al., 2016)
RL40A YEAST (Roscoe et al., 2013)
21

---

Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction
SPIKE_SARS2 (Starr al., 2020)
Q2NOS5 _HIVI (Haddox al., 2018)
TPOR.HUMAN (Bridgford et al., 2020)
LPGLSTRSG STRSG (Olson al., 2014)
G. Homologous database construction
We utilize the homologous sequence database constructed by Truong Jr & Bepler (2023) UniRef50 version 2021/03 (Suzek al., 2015). We thank the PoET authors for providing access this valuable resource. The database was constructed through an all-vs-all homology search across UniRef50 using Diamond (Buchfink al., 2015), a high-performance sequence alignment that over 100x faster than BLAST. The search was performed with the following parameters:
diamond blastp uniref50. fasta -d biamondd/niieen 0. 001 PP 96 output. .tab
-f 6 --header -k 200000 --max-hsps 1
This comprehensive search identifies putative homologs for each of the approximately 62 million sequences in UniRe(50, creating homologous clusters that the foundation for our retrieval experiments. The resulting cluster size distribution is shown in Figure G.I.
10F
106
100000
UniRef50 cluster
Figure G.I. Distribution over cluster sizes of UniRef50. Distribution of the 62 million UniRef50 clusters.
H. Software
We make our code available at https:///// / htt.aaa......e.aaaaaaaaaaaaaaaee
22


Reviews:
[{"review_id": "lqPQZoEKCJ", "reviewer": "Reviewer_2keu", "summary": "This paper presents a method for jointly training a homology retrieval model and a conditional language model for fitness prediction. The retrieval model efficiently identifies homologous sequences for a given query protein using efficient vector search.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**C1: PoET / SOTA results**\n\nOur suggested differentiable retriever approach can readily be used to augment SOTA sequence-based architectures on the ProteinGym benchmark such as PoET. The framework requires minimal changes, though the code has to be modified to allow each query in the batch to process a variable number of retrieved sequences.\nPoET leverages additional processing of sequences based on various sequence similarity cutoffs on the aligned sequences. This results in more than 30 forward passes through the model (sequences scored from N to C and C to N, with 15 different sampling hyperparameters). Note that PoET's performance when scoring only one direction with default parameters (context size of 12288 and sequence similarity cutoff of 0.95) on the validation set is 0.425. While we could integrate an alignment step in our framework, which would allow for similar inference results, we present here results of a single forward pass of the model, which we think better represents the true potential of the method.\nWe focus our analysis on the subset of 10 DMS assays from ProteinGym used as our validation set, and obtain the following results:\n\n| Assay | Spearman - vanilla PoET | Spearman - PoET w/ Protriever |\n|-|-|-|\n| BLAT_ECOLX_Jacquier_2013 | 0.677 | 0.700 |\n| CALM1_HUMAN_Weile_2017 | 0.201 | 0.208 |\n| CCDB_ECOLI_Tripathi_2016 | 0.423 | 0.422 |\n| DLG4_RAT_McLaughlin_2012 | 0.594 | 0.581 |\n| PA_I34A1_Wu_2015 | 0.527 | 0.526 |\n| Q2N0S5_9HIV1_Haddox_2018 | 0.541 | 0.518 |\n| RL40A_YEAST_Roscoe_2013 | 0.420 | 0.477 |\n| SPG1_STRSG_Olson_2014 | 0.272 | 0.216 |\n| SPIKE_SARS2_Starr_2020_binding | 0.357 | 0.356 |\n| TPOR_HUMAN_Bridgford_2020 | 0.447 | 0.395 |\n| **Average** | **0.446** | **0.440** |\n\nWe match the fitness prediction performance of PoET obtained with ColabFold (performance difference is not statistically significant; using the same bootstrap standard error (BSE) methodology from ProteinGym, BSE=0.0092)\n\n**C2: time-performance trade-off**\n\nThank you for the great suggestion.\nWe compared the timing and fitness prediction performance from 4 sequence retrieval systems (Protriever, MMseqs2 (k-mer), MMseqs2-GPU, and JackHMMER) used in conjunction with our proposed FiD architecture or PoET (as discussed in the prior response). For the MMseqs2 searches, we use the default parameters from Kallenborn & Chacon et al. (see our response on benchmarking methodology B1 to reviewer Bbfm for details). \nWe query each wild type sequence of the validation set against the same UniRef50 database as the one used in the original PoET paper (i.e., the 2023_04 checkpoint)\nWe obtain the following results on the validation set, using either the FiD model or PoET (single inference forward pass) :\n\n| Retrieval Method | Avg. Retrieval Speed (sec) | Avg. Spearman FiD | Avg. Spearman PoET |\n|-|-|-|-|\n| Protriever| 0.0046  |0.404| 0.440\n| MMseqs2| 25.9495| 0.411|0.446\n| MMseqs2-GPU| 0.7745| 0.402| 0.435\n| JackHMMER| 2645.0| 0.413| 0.449\nWe observe that Protriever achieves more than 100x speedups in retrieval, while maintaining comparable fitness prediction performance.\n\n**C3: EMDR as EM like objective**\n\nOur EMDR loss follows Sachan et al.'s \"EM-inspired\" (we use the same terminology as the authors) approach rather than classical EM. The connection to EM is:\n\n1) Implicit E-step: Latent variables are the relevant documents. We estimate their posterior distribution by computing $p_{RETR}(\\mathbf{d}_k|\\mathbf{q})$ for retrieved documents.\n\n2) Implicit M-step: We update parameters by maximizing the loss in our paper.\n\nThe loss is akin to EM's core concept of alternating between estimating latent document relevance and optimizing model parameters.\n\n**C4: Training strategies**\n\nThe ESM set strategy was initialized from the last fixed BLAST experiment checkpoint. These results are not set for comparison but to show how these three strategies of iterative training allow us to improve the performance of our retriever. Fixed BLAST sets provide the biggest boost in performance compared to the base retriever (0.347 to 0.379 for EMDR), and allows for the retriever to focus training on hard-to-distinguish sequences, as these are pulled by sequence similarity BLAST search. The second round of training with ESM sets contains easier to separate sequences and further increases performance (0.379 to 0.385). Finally, we train the reader too, end to end on the learned retriever sets, so that the reader model learns to extend to further sequences (0.385 to 0.404).", "claims_and_evidence": "The paper clearly articulates its main claim, and the empirical study aligns well with it. However, it lacks sufficient comparison with the state-of-the-art model, PoET, which is mentioned in the related work.\n\nAdditionally, while the paper asserts that the proposed method achieves strong performance with improved computational efficiency, the provided empirical studies do not convincingly support this claim. A more compelling demonstration would be a curve plotting computation time (x-axis) against performance (Spearman rank correlation, y-axis), clearly illustrating the model’s advantage in the time-performance trade-off.", "methods_and_evaluation_criteria": "The proposed method is intuitive, though the EMDR retriever loss does not appear to be an EM objective. While this may not be a major contribution of the paper, it would be helpful for the authors to provide an explanation—perhaps in the Appendix—to clarify the validity of this objective.\n\nThe evaluation criteria follow the standard setting and seem appropriate.", "theoretical_claims": "No theoretical claims are made in this paper.", "experimental_designs_or_analyses": "This paper lacks sufficient comparison with the baselines. For example, the SOTA methods in the ProteinGym are ProSST and PoET.  ProSST uses structure information, which makes it an unfair comparison with this paper. However, the PoET is closely related, and the performance is better than the performance of this model, which should be compared and discussed.\n\nThis paper lacks sufficient comparison with relevant baselines. In particular, the SOTA method PoET in ProteinGym is not adequately discussed. PoET is closely related and achieves superior performance. A thorough comparison with PoET, including a discussion of the differences and potential advantages, would strengthen the paper.", "relation_to_broader_scientific_literature": "This paper is related to scientific literature about protein fitness predictions.", "essential_references_not_discussed": "The paper appropriately cites essential references. However, as mentioned earlier, a more in-depth discussion comparing the proposed method with the SOTA, particularly PoET, should be included to provide a clearer contextual understanding of its contributions and limitations.", "other_strengths_and_weaknesses": "This paper demonstrates a fair level of originality. While the methods employed originate from prior work in other domains, such as NLP, this is the first study to apply them to protein fitness prediction, marking a novel contribution to this field.", "other_comments_or_suggestions": "NA", "questions_for_authors": "(1) Why is the proposed method not compared with PoET in Table 1?\n\n(2) In the last paragraph of Sec. 4.3, it states: 'Each method was initialized from the checkpoint of the previous method.' Does this mean that the model using the 'ESM sets' strategy was trained based on the checkpoints from the 'Fixed BLAST sets' strategy? If so, is this a fair comparison for evaluating the incremental performance of the different training strategies?\n\nI will consider raising my score if the authors address these two questions, particularly the first one.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "njy6IWqmJk", "reviewer": "Reviewer_Bbfm", "summary": "The paper proposes Protriever, an end-to-end differentiable framework for protein sequence modeling. This approach unifies two steps: protein homology retrieval and downstream modeling tasks. This is done by using a vector similarity search to retrieve homologous protein sequences. The authors train Protriever end-to-end for a protein fitness prediction task. They evaluate it on the ProteinGym benchmark for zero-shot fitness prediction. Protriever achieves performance on par with traditional MSA-based approaches. The benefit is on the speedup --- it's reported to be 30-100x faster compared to existing methods.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**C1: Speed-up: inference analysis**\n\nInstead of reporting per query times from the MMseqs-GPU paper from Kallenborn et al. we move to rigorously benchmark our method, MMseqs2 in CPU and GPU modes.\nWe randomly sample UniRef50 sequences to create 5 sets of sizes 1, 10, 100, and 1000 each. For each query size, we report the mean and standard deviation across the 5 random datasets (after taking the median over 5 repeated runs per dataset). \nWe use the same MMseqs2 parameters as described in Kallenborn & Chacon et al. to search against ~62 million UniRef50 sequences: increased sensitivity for the CPU-version (s=8.5), prefilter-mode 1, E-value cutoff 10000, and 4000 max sequences. For the GPU-version, we use flags ‘--gpu 1’ and ‘--gpu-server 1’, and the same E-value cutoff and max number of sequences. \n\nOur method and MMseqs2-GPU are run on identical hardware with a single L40S using a single thread. MMseqs2 uses 20 threads. The Protriever faiss index is pre-trained on all sequences and does not use sharding for direct single-GPU comparison.\n\nSearch time per query (in seconds) shows that Protriever remains orders of magnitude faster across query sizes\n\n| Method| N=1| N=10| N=100| N=1000|\n|:-|:-|:-|:-|:-|\n| MMseqs2 (k-mer)| 28.587 ± 5.7237 | 17.771 ± 2.9081 | 15.010 ± 1.4492 | 15.567 ± 1.0836|\n| MMseqs2-GPU| 0.637 ± 0.3079| 0.532 ± 0.1285|0.493 ± 0.0653|0.508 ± 0.0246|\n| Protriever (faiss) | 0.007 ± 0.0005| 0.005 ± 0.0001| 0.005 ± 0.0001| 0.005 ± 0.0001|\n\n**C2: Speed-up: training analysis**\n\nEmbedding the full UniRef50 database with ESM-2 (35M) takes less than 2 hours on a single GPU when using efficient flash-attention implementations. A full training and inference time comparison with other retrieved methods is practically not possible, each requiring dedicated processing steps (MMseqs2 costly database indexing, traditional homology methods compute expensive MSA profiles). Faiss's advantage is that we can export a small index that represents a “cached” result, which can easily be downloaded by end users. Fair comparisons should focus on inference time, where our method demonstrates clear superiority, as shown above. In response C1 to reviewer 2keu, we have included a table showing the performance and retrieval times of competing methods on 10 validation sets from ProteinGym. The faiss-driven retrieval time is several orders of magnitude faster, such that substituting it with MMseqs2-GPU would increase the computation time significantly.\n\n**C3: Applications benefiting from retrieval speedups**\n\nThere are many applications that could significantly benefit from the architecture we present in this work. To name a few:\n- Evaluating the effect of missense and indel mutants on the full human proteomes entails extracting 20.000+ MSAs, one for each gene, which is time and compute heavy.\n- Training structure prediction models such as AlphaFold or OpenFold. The Open Protein set that was used to train the latter is for instance composed of 16 million MSAs. Refreshing this set as underlying data sources keep being updated is intensive.\n- Rapid retrieval is similarly useful in the inference setting of structure prediction models that leverage homologous sequences, where the arguments made in Kallenborn et al. similarly apply here.\n\n**C4: Statistical significance**\n\nWe calculate non-parametric bootstrap standard errors (BSE), methodology from ProteinGym. As results in Tables 1 and 2 are reported over large number of assays they are robustly statistically significant (low MSA depth BSE<0.001). For the the standard estimate for EMDR vs. PDist as computed over the validation set only, the statistical significance is of 0.0072.\n\n**C5: Insertion/Deletion**\n\nThank you for the suggestions. We have evaluated our Fusion-in-Decoder on the set of 66 assays in the ProteinGym benchmark. We obtain a score of 0.425 for FiD evaluated with MSA retrieval and 0.422 for FiD with Protriever retrieval. We can observe once again that we match (BSE=0.01) performance of the standard MSA-retrieval approach further demonstrating the versatility of the proposed framework.\n\n**C6: Sampling Strategies at training and inference**\n\nDuring training, sequences are sampled with weight inversely proportional to Uni50 cluster sizes to avoid oversampling large clusters (will have more hits at the Uni50 level). We also sample sequences from Uni100 for each Uni50 sequence with weight inversely proportional to the size of the Uni90 clusters they belong to introduce more diversity and not overly represent larger Uni90 clusters. \nAt inference we test out different sampling strategies that combine two sampling strategies: \n- When sampling Uni100 with Uni90 sequence weights, as in training we make sure to not oversample large Uni90 clusters within each Uni50 representative.\n- When sampling based on embedding distance, we aim to sample closer sequences at Uni50 level, our retriever is trained to put closest to each query the most useful retrieved documents.", "claims_and_evidence": "The primary focus of the paper is on the speedup. The authors report the performance (which is indeed on par with the other methods). I think this focus is reasonable and, therefore, I won't comment on performance differences and focus mostly on the speedup part.\n\nComments on the speedups:\n\n- Because we're interested in the speedup, I'd like to ensure that the results of the speedup are not influenced by any other factors. What is the benchmarking methodology for this time? What hardware was used? What is the dataset size? Configuration parameters?\n\n- On Figure 2, can you report more queries than 1 and 100? Why these numbers? Can you also report uncertainty and confidence intervals?\n\n- Because the primary benefit seems to be the speedup, could you tell me and report the end-to-end performance timing for complete prediction tasks? Right now, the speedup is only for the retrieval time. I think it could be alright if the training time is much larger for your method than other methods (i.e. I think you can make the case that retrieval time is more important), but I'd like to have this analysis. So: (i) can you compare the end-to-end time?; (ii) if the end-to-end time for your application is larger with one retrieval, how many retrievals would you need to have for you to break-even time-wise?\n\n- On a more practical level, could you tell me a bit more about what applications require such speedups (i.e. are there any reports showing that speedups are issues people are trying to solve?). I'm happy if these are non-academic reports too if that helps.\n\nOther comments:\n\n- On ablation studies. The ablation studies (Table 3, Table 4) seem to have very minor differences between different values tested. Could you report any uncertainty intervals (e.g. standard deviation over different splits)? It's hard to see whether the values are significant.\n\n- The statistical significance angle is true elsewhere too. For instance, you highlight Protriever's advantage in low MSA depth regimes (0.365 vs. 0.352), but this small difference may not be statistically significant, and the paper doesn't provide statistical validation. Because you are reporting average Spearman rank correlation, I don't have good intuition whether such a difference means anything -- to me this looks like noise.\n\n- You claim that Protriever is \"architecture-agnostic,\" but this appears somewhat overstated as the flexibility to substitute different components is common to many deep learning frameworks. I appreciate the separation of retriever and reader components in your design,  but it's somewhat of a trivial point that these could be any architectures. The same is true for the task. I'd recommend either removing this framing or providing empirical evidence across multiple tasks and architectures beyond fitness prediction to make this more of an important point in the paper.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are generally appropriate. ProteinGym is a suitable and standard choice in protein fitness prediction. \n\nWhile most things are clear, here are a few questions from my end that it would be great to get answers to:\n\n- One issue is that the paper evaluates only on single amino acid substitutions and doesn't test more complex mutation scenarios (e.g. deletions, insertions).\n\n- Another issue -- repeating the comment from before -- is that the authors focus only focus on a single prediction task, so it limits your claim that this is task-agnostic (even though I understand what you mean by this).\n\n- It would be also great if you could expand your writeup to explain what you do (e.g. training procedures) instead of simply stating that you use PoET's dataset and sampling approach (which makes it hard for a reader to follow what's happening).\n\n- Also, do you have any intuition why different sampling strategies work better than others?", "theoretical_claims": "No proofs or theoretical claims.", "experimental_designs_or_analyses": "I looked at the experimental design. The issues I've noticed are already included above.", "relation_to_broader_scientific_literature": "Protriever builds upon conditional protein language models like MSA Transformer (Rao et al., 2021) and PoET (Truong Jr & Bepler, 2023), which condition on multiple related sequences. Unlike MSA Transformer, which requires pre-aligned sequences and cannot model insertions/deletions, Protriever follows PoET's approach of handling unaligned sequences, but adds the innovation of AIDO which sequences are most useful during training.\n\nProtriever adapts RAG techniques from NLP (Lewis et al., 2020) to protein modeling, similar to concurrent work like RSA (Ma et al., 2024) and AIDO.RAG (Li et al., 2024).", "essential_references_not_discussed": "N/A.", "other_strengths_and_weaknesses": "Some strengths I haven't mentioned before: \n\n(1) the framework's ability to change retrieval databses at inference time seems to be practically useful; \n\n(2) Generally, it seems to integrate multiple disciplines quite well. The EMDR and PDist training approaches represent interesting adaptations of NLP retrieval methods to the protein domain, with EMDR showing slightly better performance in the ablation studies. The Fusion-in-Decoder architecture is a sensible choice that handles the computational complexity challenges of processing multiple sequences effectively.\n\n(3) it seems to address clear problems within MSA. The end-to-end differentiable retrieval framework for protein sequences represents a novel technical contribution that hasn't been previously explored in this domain.\n\nThe drawbacks are largely those that relate to the questions I had asked previously.", "other_comments_or_suggestions": "N/A\n\n## Post-rebuttal update: Increased 3 -> 4 based on the reply.", "questions_for_authors": "Asked above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9IX9WfgAJ9", "reviewer": "Reviewer_cdq5", "summary": "Protriever is an end-to-end differentiable framework for augmenting performance on downstream applications (e.g., fitness prediction)  of language models using vector-based retrieval. The model consists of 2 trainable components, the retriever and the reader, and one static vector index. The retriever learns to select from the index a set of support sequences for a query sequence while the reader learns to reconstruct the query from only the support sequences (by autoregressive decoding). In so doing, the reader learns a joint distribution over query and support sequences. After training, protriever is applied to a fitness prediction task.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**C1: EMDR optimum**\n \nThe EMDR loss function is defined as:  \n\n$$ \\mathcal{L}_{EMDR} = -\\log \\left[  \\sum_k p^{LM} ( \\mathbf{q} \\vert \\mathbf{d}_k ) p^{RETR}  ( \\mathbf{d}_k \\vert \\mathbf{q}) \\right] $$\n\n(We are experiencing issues with longer equations in OpenReview's Markdown, where subscripts appear to break the rendering - so we have resorted to put _LM_ and _RETR_ as superscripts.)\n\nMinimizing this loss is equivalent to maximizing the weighted sum inside the logarithm. Given that $p_{RETR}$ must form a valid probability distribution (summing to 1), this becomes a constrained optimization problem.  \n\nThe solution to such a problem is straightforward: to maximize a weighted sum $\\sum_k w_k p_k$ where weights $w_k$ are fixed and $\\sum_k p_k = 1$, the optimal strategy is to assign $p_j = 1$ to the term with the highest weight $w_j$, and $p_k = 0$ to all others.  \n\nIn our case, the weights are the language model probabilities $p_{LM}(\\mathbf{q} | \\mathbf{d}_k)$, so the optimal retriever distribution will place all probability mass on the document that yields the highest language model score.  \n\nThis is intuitive when viewing the retriever as allocating a fixed \"budget\" of probability mass to maximize the weighted sum — the best strategy is always to invest the entire budget in the option with the highest return.\n\n**C2: Additional References**\n\nThank you for the suggestion — both are relevant and will be included in the revision.\n\n**C3: Validation Set**\n\nFor the validation set, we have used the same set as previously used in Tranception (Notin et al) and PoET (Truong et al). It is composed of the following assays:\n- BLAT_ECOLX (Jacquier et al., 2013)\n- CALM1_HUMAN (Weile et al., 2017)\n- CCDB_ECOLI (Tripathi et al., 2016)\n- DLG4_RAT (McLaughlin et al., 2012)\n- PA_I34A1 (Wu et al., 2015)\n- RL40A_YEAST (Roscoe et al., 2013)\n- SPIKE_SARS2 (Starr et al., 2020)\n- TPOR_HUMAN (Bridgford et al., 2020)\n- Q2N0S5_9HIV1 (Haddox et al., 2018)\n- SPG1_STRSG (Olson et al., 2014)\n\nWe will include the above clarification in the revision.\n\n**C4: Performance impact and breadth of tasks supported**\n\nAs noted by the two other reviewers, the core practical benefit from our suggested approach lies in the significant speedups over standard retrieval methods such as multiple sequence alignments (MSA), without degrading the fitness prediction performance. We point the reviewer to the response to C2 from reviewer 2keu and C1-C2 from reviewer Bbfm for a more detailed discussion of our contributions there. Note that MSAs have been a cornerstone of computational biology for several decades and, as such, they have been extensively optimized to support tasks like the ones we focus on in this work. Thus, the types of improvements we present in this work are both non-trivial and of significant importance in practice, where these speedups matter (C3, Bbfm)\n\nWe hope this addresses the concerns of the reviewer in terms of performance improvements. As for the scope of applications our paper covers, we would like to emphasize that ProteinGym encompasses 200+ assays across a wide range of protein functions (e.g,. thermostability, binding, catalytic activity, protein expression) and that, as such, our current experimental scope is already fairly broad and diverse. As noted by the reviewer, there are other protein-related applications where homology is key (e.g., function prediction, tertiary structure prediction). While an adaptation of our proposed method to these other tasks is beyond the scope of this paper, we argue that this potentiality is yet another strength of our proposed model and believe the ideas presented here will be conducive to many subsequent works that explore the benefits of end-to-end retrieval in protein modeling.", "claims_and_evidence": "The claims/result on the architectural choices are well-justified and convincing. The claim that the training procedure improves fitness prediction is supported by Tables 1 and 2, which seem to show a significant improvement over the tested baselines in only certain settings.", "methods_and_evaluation_criteria": "The evaluation criteria is sound -- the ProteinGym dataset is a well known benchmark for fitness prediction.", "theoretical_claims": "I would like for the authors to provide a reference or brief justification of the claim in sec 3.2 line 245 that the theoretical optimum of the EMDR loss is \"a degenerate distribution that assigns all probability mass to the single sequence maximizing the language model’s likelihood of generating the correct output\".", "experimental_designs_or_analyses": "I did check the soundness/validity of the experiments.", "relation_to_broader_scientific_literature": "The work relates to ongoing advancements in sequence retrieval for protein property prediction, such as ProtEx [Shaw, et al '24] and builds off seminal work on retrieval augmented generation (RAG, REALM) referenced in the text.", "essential_references_not_discussed": "- Alongside the reference to MSAGPT, the authors may also refer to EvoDiff [Alamdari, et al '23].\n- On the subject of structure-aware embeddings models (sec 5, line 410), the authors may also refer to TM-Veec [Hamamsy, et al '24, published in Nat. Bio].", "other_strengths_and_weaknesses": "Strengths\n\nThe paper is well-written, clear, and the architectural choices are well-motivated and justified.\n\nWeaknesses\n\nThe article emphasizes architectural decisions but the experimental results are not convincing. The model does not improve performance significantly or across-the-board. The application to fitness prediction has a well-defined benchmark, but the results do not seem to demonstrate any marked effectiveness of the model for this task.\n\nI am not sure what Table 3 tells us, since the Spearman values are generally close. Which loss function was used in the end and why?", "other_comments_or_suggestions": "- Which assays were used during the ProteinGym evaluation?", "questions_for_authors": "There are many other potential applications of this model besides fitness prediction where there doesn't seem to be a clear demonstration of improved performance. I am curious whether the authors are willing to explore any other applications which could separately test the retriever and reader's capabilities. For example, both the reader and retriever can be used for function prediction (by means of homology transfer in the latter).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces Protriever, an end-to-end differentiable framework for protein sequence modeling that unifies protein homology retrieval and fitness prediction. Protriever employs a vector similarity search to efficiently retrieve homologous sequences and trains the retriever and reader components jointly for fitness prediction. The framework achieves performance comparable to traditional MSA-based methods while being 30-100x faster, as demonstrated on the ProteinGym benchmark for zero-shot fitness prediction. Although heavily borrowing from prior work in NLP, this is the first work to apply joint retrieval and conditional language modeling to protein fitness prediction, which is significantly faster than traditional MSA-based approaches.

All reviewers unanimously agree to accept the paper after the rebuttal phase, so is the final decision.
Decision: Accept (poster)