Training on test proteins improves fitness, structure, and function prediction

W1. The method is motivated by models with "Y-shaped" architectures in which there is both an MLM head and a fitness prediction head. However, the authors evaluate on the zero-shot task in ProteinGym, where models use masked token probabilities as proxies for fitness predictions and therefore do not use a Y-shaped architecture. This is confusing, since ProteinGym also contains fine-tuning tasks that require the Y-shaped architectures that motivate their method. The paper would be improved by evaluating on the fine-tuning tasks in ProteinGym or by more clearly explaining why the zero-shot task is more appropriate.

We have clarified that the zero-shot setup also follows the Y-shaped architecture. Please note that in the zero-shot setup, the Y-shaped architecture is given by the language model backbone $f$ (e.g., ESM2 Transformer encoder), the pre-training head $g$ given by the MLM head, and the downstream-task head $h$ given by the standardly used log odds ratio. This may not be obvious to realize since, in this case, $h$ contains $g$. We have clarified this aspect in Section 3.2.

We have clarified the choice of the zero-shot ProteinGym benchmark as the best established benchmark for fitness prediction. The goal of our paper is to demonstrate the potential of TTT across well-established problems and benchmarks. Therefore, we chose the best available benchmark for fitness prediction: ProteinGym, specifically its zero-shot setup. It also allows us to evaluate TTT in a different setting compared to the structure and function prediction tasks, where downstream task heads $h$ are fully supervised. We have updated the Evaluation Setup paragraph of Section 4.1 to incorporate this clarifications.

We have clarified why we did not employ the supervised ProteinGym benchmark. To the best of our knowledge, unlike the zero-shot ProteinGym benchmark, the supervised ProteinGym benchmark has not yet been widely adopted and currently features only methods from a single publication, ProteinNPT (Notin et al., 2023). According to the official ProteinGym GitHub page, the supervised benchmark is not fully integrated into the ProteinGym codebase and is only accessible through the ProteinNPT codebase. We have now explained it in the ProteinGym paragraph of Section A.1.1 in the Appendix. When experimenting with the ProteinNPT codebase, we encountered computational bottlenecks while reproducing the results of the baseline methods. Specifically, each model needs to be trained on every assay, requiring approximately 24 hours per assay. Unfortunately, after carefully considering the suggestion to evaluate TTT on this benchmark, we found this task infeasible to complete within the rebuttal period.

Conceptually, however, we do not foresee any obstacles preventing TTT from improving the state-of-the-art ProteinNPT model. This model follows the Y-shaped architecture and leverages masked modeling as extensively as the other methods considered in our work. It does so in two ways: first, by utilizing MSA Transformer embeddings and zero-shot predictions as inputs, and second, through the non-parametric transformer approach (NPT), which masks multiple sequences in a batch. To strengthen our argument, we provide results from evaluating MSA Transformer + TTT in a zero-shot setting (please see our response to W2).

Since ProteinNPT uses MSA Transformer zero-shot predictions and embeddings as inputs, these improvements suggest that TTT could also benefit ProteinNPT when applied. Additionally, TTT can be directly applied to the NPT batch-wide masking step itself, which is expected to further enhance its performance.

W1. TTT Method Selection: The rationale for selecting the current TTT methods is unclear. Provide a clear justification for the choice.

We have modified the opening paragraph of Section 2 to clarify that the section provides the rationale, feasibility, and broad applicability of the TTT method. We have also improved the overall clarity of our text to better highlight the justification of our method.

W2. Performance Linkage: The objectives should be linked to improved biological performance. Add relevant goals to achieve desired outcomes.

We have extended the motivation for applying TTT to protein structure prediction by discussing its biological implications in Section 4.2:

• “Arguably, one of the most remarkable applications of machine learning in the life sciences has been in protein folding (Jumper et al., 2021; Lin et al., 2023; Abramson et al., 2024), paving the way for numerous advances in the understanding of biology (Yang et al., 2023; Akdel et al., 2022; Barrio-Hernandez et al., 2023).”

The objectives for developing TTT for proteins are currently outlined in the Introduction, while the applications of TTT to protein fitness, structure, and function prediction are linked to improved biological performance in the opening paragraphs of Sections 4.1, 4.2, and 4.3, respectively:

• “Bridging the gap between broad, dataset-wide optimizations and the precision required for studying single proteins in practical applications remains a critical challenge in integrating machine learning into biological research (Sapoval et al., 2022). … Our method enables adapting protein models to one protein at a time, on the fly, and without the need for additional data.”
• “Protein fitness refers to the ability of a protein to efficiently perform its biological function, which is determined by its structure, stability, and interactions with other molecules. … In this paper, we demonstrate that applying test-time training (TTT) to representative models, such as ESM2 (Lin et al., 2023) and SaProt (Su et al., 2023), enhances their protein fitness prediction capabilities.”
• “Protein function prediction is essential for understanding biological processes and guiding bioengineering efforts. … Therefore, enhancing generalization through TTT is particularly relevant in this context.”

We are open to extending the discussion of our goals to include any additional specific objectives that might be suggested.

W3. Efficiency: Reducing perplexity is insufficient for justifying large model use. Considering smaller discriminators or LLM query filtering for efficiency is the same.

It is a valid point that reducing perplexity is insufficient for justifying large model use. Please note that the goal of our method is to adapt a given protein LLM, regardless of its size, to a specific protein sequence, with perplexity reduction as the means to achieve this. We focus on large protein language models to demonstrate the applicability of our method to state-of-the-art models, which are typically large, rather than limiting results to smaller, proof-of-concept models. We have clarified this aspect in the paragraph “Fine-tuning large models” in Section 3.1:

• “Since state-of-the-art models for many protein-oriented tasks are typically large, with up to billions of parameters, our aim presents two key challenges… To address the first challenge, we perform forward and backward passes through a small number of training examples and accumulate gradients… We address the second challenge by employing low-rank adaptation…”

W4. Resource Management: Finetuning large models is resource-intensive. Assess cost-effectiveness and seek more efficient training strategies.

Thank you for pointing this out. We have added a new Figure 9 in the Appendix highlighting the cost-effectiveness of our TTT approach using low-rank adaptation (LoRA) and gradient accumulation. We have referenced it in the main text:

• “Figure 9 in Appendix B shows that ESMFold + TTT maintains computational efficiency comparable to ESMFold while being orders of magnitude faster than AlphaFold2.”

W1. My main concern is the proper estimation of the number of test-time training steps. Such fine-tunings have been studied frequently and are prevalant in other areas, such as few-shot learning, meta-learning, domain adaptations, etc. Most, if not always, (1) the size of the fine-tuned parameters, and (2) the number of training steps can play a significant role in the outcome.

While the paper is not devoid of studies regarding these effects, I don't feel they are complete enough. Once would expect a clear degradation in performance with large-enough step sizes and model-complexity. Some intuitive indications of when and how TTT should be proposed, studied, and tested. There is almost no theoretical component to the paper, which is not a deal breaker for this paper, but it may assist in the correct choice of such hyper-parameters.

We thank the reviewer for the thoughtful comments, which we address below.

The number of training steps

For each task, we used held-out validation data to fine-tune three primary hyperparameters: learning rate, batch size (or the number of gradient accumulation steps), and the number of TTT steps. The same final task-specific set of hyperparameters was applied to each target test protein, and the overall results were reported. As shown in Table 4 in the Appendix, we experimented with different learning rates and batch sizes (“Grad. acc. steps”) while keeping the number of TTT steps fixed. The rationale for fixing the number of steps is illustrated in Figure 10, which demonstrates that—despite having three hyperparameters—there are effectively only two degrees of freedom controlling the performance of the model. In other words, by keeping the number of steps constant, the expected performance can be controlled by adjusting the learning rate and the batch size. We have clarified this observation in the caption of Figure 10. We have also clarified the caption of Table 4 to highlight that, for this reason, we used 30 TTT steps in all our experiments. The only exception was ESM3 + TTT, where the number of steps was inconsistently set to 50 during initial experiments with different models and tasks conducted in parallel before standardizing the number of steps to 30.

For the task of protein structure prediction, we determined the optimal number of steps for each sequence using a well-established measure of model confidence: the pLDDT confidence score of each residue predicted in the structure, averaged across the sequence. We have added a new Figure 11 in the Appendix, demonstrating that, with the final learning rate and batch size, the number of steps could be reduced to improve computational efficiency while still achieving a significant improvement in performance. However, as shown in the newly added Figure 9 in the Appendix, our approach remains efficient even with 30 steps.

In future work, we aim to explore other heuristics for estimating the number of steps, focusing on methods that assess the quality of protein representations from a protein language model and can be applicable to any downstream task, such as linear probing of protein contact maps (Rives et al., 2020) or the recently proposed categorical Jacobian (Zhang et al., 2024).

The size of the fine-tuned parameters

Our work suggests the following relationships between the number of fine-tuned parameters and the performance of TTT:

• TTT benefits smaller models more than larger models. This is supported by Table 5 and Table 6 in the Appendix, which compare the performance gains from applying TTT to ESM2 (35M) and ESM2 (650M), as well as SaProt (35M) and SaProt (650M). Across the two datasets and two models, TTT consistently results in larger improvements on smaller models.
• LoRA is only beneficial when the backward pass through the full model does not fit into GPU memory. We have extended Table 4 in the Appendix to show that we experimented with applying LoRA to ESM2 (35M) and ESM2 (650M), both of which can be fine-tuned without it. In line with the original LoRA paper (Hu et al., 2021), we did not observe any improvements in downstream task metrics when using LoRA in these cases.
• Fine-tuning all layers outperforms fine-tuning a subset of layers. When predicting protein structures, the ESMFold head $h$ extracts information from the ESM2 backbone $f$ by computing a learnable linear combination of embeddings across all layers. During initial experiments with ESMFold + TTT, we observed that 92% of the weights are attributed to the final layer and 5% to the second layer. Motivated by these findings, we experimented with fine-tuning only these two layers or a subset of the last layers in various combinations. However, the performance never surpassed the performance achieved when fine-tuning all layers.

I would like to thank the authors for addressing my comments and keep my score as it was.

W1. As acknowledged by the authors, the TTT technique has been used in other domains, however this study applies it for the first time to the computational problems in protein domain.

Thank you for this observation. We acknowledge that the TTT technique has been used in other domains, and we cite the relevant papers in the Background and Related Work sections. Building on ideas from other areas of machine learning has led to powerful methods in the past. Examples include protein language models derived from natural language models or protein structure diffusion models inspired by generative models in computer vision. However, we would like to point out that in order to extend TTT to the protein domain we had to address several challenges, as described below.

• In computer vision, TTT primarily addresses explicitly known distribution shifts (e.g., blurred or rotated images) (Sun et al., 2020, Gandelsman et al., 2022), whereas our work focuses on unknown distribution shifts, which are common when applying machine learning to understudied proteins.
• While the reasons for TTT’s effectiveness are often unclear in other domains (e.g., in computer vision), we provide a possible explanation in the domain of proteins through perplexity minimization (Section 3.3).
• Making TTT applicable to a single new unseen target protein on-the-fly required important design choices not addressed in the original TTT papers in computer vision. This includes the use of LoRA (Section 3.1), gradient accumulation (Section 3.1), and the construction of a new dataset from MaveDB to rigorously optimize hyperparameters (Section A.1.1).

To the best of our knowledge, our work represents the first application of the TTT approach to biological data, where adaptation to individual examples is particularly important, as discussed in the first paragraph of the Introduction and also recognized by Reviewer dg8u.

W2. One of the motivations of this work is to get rid of additional data collection (MSA) for a given test protein, however the experiments do not show this in a comparable setting. I may be wrong, so the author’s clarification is appreciated. Please see the questions for more detailed on this.

We would like to kindly clarify that the motivation of our work is not to remove the need for additional data collection but rather to provide a universal customization method that does not require additional data. This makes our method orthogonal to the availability of additional data, such as MSAs, allowing it to be combined with such data, if available. We demonstrate the possibility of combining our approach with MSA in new experiments described in the following comments.

W3. (i) It is shown that the performance is improved by applying TTT to foundation models, e.g. ESM2. However, there is no comparable setting in which the foundation model is fine-tuned on the sequences homologous to the test protein. Therefore, it is not clear when one should choose to go with TTT instead of MSA, i.e., incorporating evolutionary features.

Please kindly note that TTT and evolutionary features are not competing approaches. In this work we have developed a test-time training approach for fine-tuning a model on a single sequence using masked modeling. A similar approach can also be applied to fine-tune a model on the whole MSA using masked modeling (please see response to W4 for the first results in this direction). In the current paper we focused on the single sequence setup because it is conceptually easier and, therefore, more suitable for our work in test-time training for proteins. In many scenarios, using a single sequence is also beneficial compared to using the whole MSA. For example, single-sequence-based test-time training is beneficial in the following cases:

• The computational efficiency of the predictive method is important. For example, ESMFold+TTT presented in our paper can be applied on large scale to metagenomic data (e.g., to extend ESM Atlas; Lin et al., 2023), while relying on MSA would make a method an order of magnitude slower and not applicable on the large scale (please see newly added Figure 9 in the Appendix).
• MSA is not available or very poor. Table 5 in the Appendix demonstrates that TTT is most beneficial for proteins that have poor MSA. This highlights that TTT developed in this work is especially useful where MSA-based methods fail.

Dear Reviewers, we appreciate your time and thoughtful feedback. As the discussion period is nearing its conclusion, we kindly invite you to review our responses at your convenience (if you have not done so already). If there are any points that remain unclear or require further elaboration, we would be more than happy to provide additional clarifications promptly. We also hope that our responses have addressed your concerns. Thank you again for your time, effort, pointers to related work and invaluable insights.