Protriever: End-to-End Differentiable Protein Homology Search for Fitness Prediction

C1: EMDR optimum

The EMDR loss function is defined as:

$\mathcal{L}_{EMDR} = -\log \left[ \sum_k p^{LM} ( \mathbf{q} \vert \mathbf{d}_k ) p^{RETR} ( \mathbf{d}_k \vert \mathbf{q}) \right]$

(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.)

Minimizing 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.

The 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.

In 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.

This 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.

C2: Additional References

Thank you for the suggestion — both are relevant and will be included in the revision.

C3: Validation Set

For 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:

• BLAT_ECOLX (Jacquier et al., 2013)
• CALM1_HUMAN (Weile et al., 2017)
• CCDB_ECOLI (Tripathi et al., 2016)
• DLG4_RAT (McLaughlin et al., 2012)
• PA_I34A1 (Wu et al., 2015)
• RL40A_YEAST (Roscoe et al., 2013)
• SPIKE_SARS2 (Starr et al., 2020)
• TPOR_HUMAN (Bridgford et al., 2020)
• Q2N0S5_9HIV1 (Haddox et al., 2018)
• SPG1_STRSG (Olson et al., 2014)

We will include the above clarification in the revision.

C4: Performance impact and breadth of tasks supported

As 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)

We 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.

C1: Speed-up: inference analysis

Instead 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. We 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). We 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.

Our 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.

Search time per query (in seconds) shows that Protriever remains orders of magnitude faster across query sizes

C2: Speed-up: training analysis

Embedding 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.

C3: Applications benefiting from retrieval speedups

There are many applications that could significantly benefit from the architecture we present in this work. To name a few:

• 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.
• 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.
• 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.

C4: Statistical significance

We 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.

C5: Insertion/Deletion

Thank 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.

C6: Sampling Strategies at training and inference

During 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. At inference we test out different sampling strategies that combine two sampling strategies:

• When sampling Uni100 with Uni90 sequence weights, as in training we make sure to not oversample large Uni90 clusters within each Uni50 representative.
• 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.

C1: PoET / SOTA results

Our 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. PoET 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. We focus our analysis on the subset of 10 DMS assays from ProteinGym used as our validation set, and obtain the following results:

We 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)

C2: time-performance trade-off

Thank you for the great suggestion. We 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). We 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) We obtain the following results on the validation set, using either the FiD model or PoET (single inference forward pass) :

C3: EMDR as EM like objective

Our 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:

1. 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.

2. Implicit M-step: We update parameters by maximizing the loss in our paper.

The loss is akin to EM's core concept of alternating between estimating latent document relevance and optimizing model parameters.

C4: Training strategies

The 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).