Training on test proteins improves fitness, structure, and function prediction (2411.02109v1)

Abstract: Data scarcity and distribution shifts often hinder the ability of machine learning models to generalize when applied to proteins and other biological data. Self-supervised pre-training on large datasets is a common method to enhance generalization. However, striving to perform well on all possible proteins can limit model's capacity to excel on any specific one, even though practitioners are often most interested in accurate predictions for the individual protein they study. To address this limitation, we propose an orthogonal approach to achieve generalization. Building on the prevalence of self-supervised pre-training, we introduce a method for self-supervised fine-tuning at test time, allowing models to adapt to the test protein of interest on the fly and without requiring any additional data. We study our test-time training (TTT) method through the lens of perplexity minimization and show that it consistently enhances generalization across different models, their scales, and datasets. Notably, our method leads to new state-of-the-art results on the standard benchmark for protein fitness prediction, improves protein structure prediction for challenging targets, and enhances function prediction accuracy.

• The paper introduces Test-Time Training, a method that fine-tunes pre-trained protein models on individual proteins to boost prediction performance.
• It employs a self-supervised masked language modeling approach to adapt the backbone model during test time, reducing perplexity.
• The technique significantly improves predictions in fitness, structure, and function, especially for proteins with limited training data.

Evaluation of Test-Time Training in Protein Prediction Models

The paper presents a novel approach to improving protein prediction tasks by applying Test-Time Training (TTT). This technique involves the self-supervised fine-tuning of protein models during test time, particularly focusing on the single protein of interest. The method aims to enhance model generalization, yielding state-of-the-art results across various protein-related predictions, such as fitness, structure, and function.

Methodology

Traditional models for protein predictions, while powerful, often struggle with specificity for individual proteins, mainly due to data scarcity and distribution shifts in large datasets. The paper proposes a shift from the general approach, using TTT to adapt pre-trained protein models to a specific protein at test time, thus bridging the gap between broad dataset-wide optimizations and precise, protein-specific insights.

TTT leverages the prevalent use of masked LLMing (MLM) in protein machine learning, employing it as the objective for self-supervised fine-tuning. Specifically, during TTT, the backbone of the model (f) is adapted to reduce perplexity on the given protein sequence while the task-specific head (h) remains fixed, thus maintaining task-specific priors and leveraging improved representations learned by f.

Results

The application of TTT to various models demonstrated consistent improvement across multiple protein-related tasks:

1. Protein Fitness Prediction: The application of TTT to models like ESM2 and SaProt not only improved their performance on datasets like ProteinGym and MaveDB but also surpassed existing benchmarks, notably in phenotypes such as organismal fitness and binding. The improvement was particularly significant on proteins with low representation in training data, highlighting TTT's utility in scenarios of data scarcity.
2. Protein Structure Prediction: Using datasets like CAMEO, the paper shows that models such as ESMFold and ESM3 enhanced their performance significantly with TTT, outperforming baselines that applied different approaches like masked predictions or chain-of-thought decoding.
3. Protein Function Prediction: The method improved classification accuracy in tasks involving terpene synthase substrates and subcellular localization, emphasizing the broad applicability of TTT across different classification settings.

Theoretical and Practical Implications

The paper establishes a link between minimizing perplexity on a single protein and improved downstream performance. This insight not only helps explain TTT's effectiveness but also informs future work in applying TTT to other domains. Practically, TTT’s ability to fine-tune complex models on the fly can be invaluable in real-world applications where specific proteins of interest must be analyzed without abundant related data available.

Future Directions

The research opens up several avenues for future work, such as exploring deeper understanding of TTT’s success and failure modes and extending these techniques to more complex foundation models. Additionally, exploring adaptation methods like domain adaptation and adaptive risk minimization could further enhance protein model generalization and adaptation capabilities.

In summary, the paper makes a strong case for TTT in enhancing machine learning predictions specific to individual proteins, addressing the boundaries of current model generalizations, and setting a research path toward more targeted and efficient protein analysis methodologies.