Metalic: Meta-Learning In-Context with Protein Language Models

## Abstract

Abstract Predicting the biophysical and functional properties of proteins is essential for in silico protein design.
Machine learning has emerged as a promising technique for such prediction tasks.
However, the relative scarcity of in vitro annotations means that these models often have little, or no, specific data on the desired fitness prediction task.
As a result of limited data, protein language models (PLMs) are typically trained on general protein sequence modeling tasks, and then fine-tuned, or applied zero-shot, to protein fitness prediction.
When no task data is available, the models make strong assumptions about the correlation between the protein sequence likelihood and fitness scores.
In contrast, we propose meta-learning over a distribution of standard fitness prediction tasks, and demonstrate positive transfer to unseen fitness prediction tasks.
Our method, called Metalic (Meta-Learning In-Context), uses in-context learning and fine-tuning, when data is available, to adapt to new tasks.
Crucially, fine-tuning enables considerable generalization, even though it is not accounted for during meta-training.
Our fine-tuned models achieve strong results with 18 times fewer parameters than state-of-the-art models.
Moreover, our method sets a new state-of-the-art in low-data settings on ProteinGym, an established fitness-prediction benchmark.
Due to data scarcity, we believe meta-learning will play a pivotal role in advancing protein engineering.

## 1 Introduction

The accurate prediction of functional and biophysical properties of proteins, such as stability or binding affinity, is a critical challenge in the physical sciences with far-reaching implications for drug discovery, medical research, and agriculture.
Protein design seeks to optimize one, or a collection of these properties, collectively referred to as fitness.
While protein fitness can be measured in vitro, the process is laborious and time-consuming.
Consequently, machine learning models have emerged as a powerful tool to accelerate protein design by predicting fitness directly from amino acid sequences in silico.
However, due to the complex, high-dimensional relationship between protein sequences and fitness, and the limited availability of high-quality data, accurate fitness prediction is a challenge.

Proteins can be encoded as sequences of characters representing amino acids, making protein language models (PLMs) effective for modeling them (Madani et al., ; Rives et al., ; Rao et al., ; Lin et al., ; Notin et al., ; Truong Jr & Bepler, ).
By predicting masked amino acids, or subsequent amino acids, over known proteins at scale, PLMs can capture much of the structure and resultant properties deriving from the amino acid sequence.
While PLMs are not directly trained to predict fitness, they are trained to model the likelihood of naturally occurring proteins, which has been found to correlate strongly with their fitness (Truong Jr & Bepler, ).
In practice, large PLMs are fine-tuned on downstream protein fitness data, often with linear regression heads added (Krause et al., ).
While this paradigm is highly effective, when there is severely limited data, or no data, PLMs provide limited utility.
With limited data, learning the correct regression is difficult, even with informative pre-trained representations.
With no data, it must be assumed that the protein fitness is solely a function of protein likelihood (Truong Jr & Bepler, ; Hawkins-Hooker et al., ), and for masked language models, it is also assumed that each amino acid contributes independently to fitness (Meier et al., ; Notin et al., ).

Figure: Figure 1: An overview of meta-learning for protein fitness prediction. PLMs are trained over massive quantities of unlabeled data. Using meta-learning, we also train over a smaller quantity of labelled fitness data. Using this extra data is critical given limited data for fine-tuning at test time.
Refer to caption: x1.png

While limited or no data may be available for specific protein fitness prediction tasks, there are often other tasks that can serve as a valuable source of guidance.
For example, rather than assuming that fitness is solely a function of protein likelihood, when no data is available for a given task, we can use data from other tasks to learn the relationship between PLM embeddings and fitness.
Due to advances in high-throughput assays, such as deep mutational scanning Fowler & Fields (), data from other tasks is available to learn this relationship.
One example is ProteinGym (Notin et al., ), which compiles over one hundred distinct fitness prediction tasks.
Training over such a distribution of related tasks, to reason about new tasks, is referred to as meta-learning (Huisman et al., ; Hospedales et al., ; Beck et al., ).

To address the challenge of limited data in protein fitness prediction, we propose Metalic, integrating in-context meta-learning, protein language models, and fine-tuning.
Metalic builds on top of PLM embeddings and fine-tuning methods, but additionally, unlike baseline methods, adds a meta-learning phase over other protein prediction tasks.
While some PLMs use in-context data (Rao et al., ; Notin et al., ; ; Truong Jr & Bepler, ), these methods do not meta-learn how to use their limited context for protein fitness prediction.
The meta-learning phase is depicted in Fig. 1, and is critical to learning the relationship between PLM embeddings and fitness, given the limited fitness data available in each task at test time.
While some meta-learning methods combine in-context learning and fine-tuning (Rusu et al., ; Vuorio et al., ), they do so by using computationally expensive higher-order gradients to account for fine-tuning during meta-learning.
Crucially, Metalic leverages in-context meta-learning, and then subsequent fine-tuning, without accounting for fine-tuning during meta-learning.
This novel combination is particularly computationally efficient, and enables Metalic to outperform more complicated methods for meta-learning.

We present Metalic, an in-context meta-learning approach to tackle the problem of protein fitness prediction in low-data settings, and make the following contributions:

- •
We introduce a method that efficiently combines in-context meta-learning with PLMs and fine-tuning for protein fitness prediction.
- •
We advance state-of-the-art (SOTA) for zero-shot protein fitness prediction on the ProteinGym benchmark (Notin et al., ).
- •
We attain strong performance for few-shot fitness prediction with 18 times fewer parameters.
- •
We ablate each component of our method to demonstrate the contributions of each part and underscore their necessity.
- •
We empirically validate the superiority of our method to alternative forms of meta-learning.

## 2 Related Work

Meta-Learning Meta-learning aims to create a sample-efficient learning algorithm by training over a distribution of tasks.
The goal is to learn algorithms such that they can rapidly adapt to new tasks during inference.
This inference-time adaptation is often called the *inner loop*, for which there are two primary forms found in the literature: gradient-based meta-learning (Finn et al., ; Zintgraf et al., ) and in-context meta-learning (Mishra et al., ; Nguyen & Grover, ).

Gradient-based algorithms and in-context algorithms differ both in their computational efficiency and capacity for out-of-distribution generalization.
Gradient-based approaches explicitly adapt model parameters within in the inner loop using standard gradient-based learning.
Commonly, the model initialization are the only parameters learned over the task distribution.
The meta-training process explicitly computes gradients through this adaptation process, effectively learning a parameter initialization that can be adapted to new tasks with only a few gradient steps (Finn et al., ; Zintgraf et al., ; Vuorio et al., ).
However, this comes with considerable computational overhead – in particular when calculating meta-gradients through the inner loop – which makes gradient-based meta-learning less suitable for large models.
Alternatively, in-context meta-learning adapts by conditioning on a task-specific dataset in context, i.e., conditioning on the data points over which gradient-based approaches would train.
Typically, this is achieved by modelling the inner loop with a sequence model that can explicitly condition on the provided data (Santoro et al., ; Mishra et al., ; Beck et al., ; ).
Such methods are typically more sample- and compute-efficient than gradient-based methods, but can perform worse when provided with out-of-distribution tasks given the lack of explicit gradient-based learning in the inner-loop (Beck et al., ).

In Metalic, we only train for in-context meta-learning, but find this is still compatible with task-specific fine-tuning.
In the meta-reinforcement learning setting, this combination has been shown to be possible by increasing task-specific data for fine-tuning (Xiong et al., ).
In contrast, we evaluate in the supervised setting and do not give increased data at inference time.
While prior work has combined gradient-based and in-context meta-learning (Rusu et al., ; Vuorio et al., ), these works compute expensive meta-gradients to learn how to account for fine-tuning.
Despite not explicitly meta-learning gradient-based adaptation, we find in-context meta-learning alone provides a strong foundation for subsequent fine-tuning and that both aspects are critical for high performance.

Likelihood-Based Fitness Prediction with PLMs Leveraging pre-trained PLMs is standard practice in protein fitness prediction (Rives et al., ; Notin et al., ; Rao et al., ; Truong Jr & Bepler, ).
In the few-shot setting, PLMs intended for sequence generation are repurposed by fine-tuning for protein fitness prediction (Rives et al., ).
In the zero-shot setting, it is assumed that the fitness correlates with the likelihood of the proteins associated sequence of amino acids, as predicted by a PLM (Meier et al., ; Truong Jr & Bepler, ).
Furthermore, if using a masked PLM, it is often assumed that each amino acid contributes independently to the fitness (Meier et al., ).
In this work, we likewise leverage PLMs for protein fitness prediction.
However, in contrast, we make use of additional data in the form of additional fitness prediction tasks on other proteins.
Specifically, we meta-learn how to use a PLM for protein fitness prediction, rather than relying on assumptions.
Only after meta-learning, do we fine-tune our model, as depicted in Fig. 1.
Meta-learning across tasks lets us avoid restrictive model constraints and achieve SOTA performance.
We will show that in-context meta-learning is necessary to achieve strong results in low-data settings.

In-Context PLMs We build upon existing PLMs that make use of in-context data for protein fitness prediction (Notin et al., ; Truong Jr & Bepler, ; Notin et al., ; Rao et al., ).
However, these methods do not meta-learn how to make use of their context.
These methods either learn to use the context only for protein language modelling, and then assume that the likelihood from the generative model correlates with fitness (Truong Jr & Bepler, ; Notin et al., ; Rao et al., ), or they use the context for protein fitness prediction, but not by meta-learning over protein tasks (Notin et al., ).
Of these ProteinNPT (Notin et al., ) is the most related to our method, since we use the same attention architecture to condition on fitness information about related proteins in-context, and it uses gradient steps to fine-tune to the target task.
In comparison, our method meta-learns over many tasks how to make use of the fitness information, which we find to be critical (Section 4).
Additionally, our method is the first to use the aforementioned procedure to allow fine-tuning and in-context conditioning on the very same context at inference time.

Figure: (a) Meta-training
Refer to caption: x2.png

## 3 Methods

### 3.1 Problem Setting

We consider a fitness prediction task, $\mathcal{T}$, to be defined by a dataset of the form ${\mathcal{D}_{\mathcal{T}}=\{(x_{i},y_{i}{\equiv}f_{\mathcal{T}}(x_{i}))\}_{i=
1}^{N}}$, where $x_{i}$ is a sequence of amino acids, and $y_{i}\in\mathbb{R}$ is the associated scalar fitness value assigned by the (unknown) underlying fitness function $f_{\mathcal{T}}$.
In the few-shot setting, the task-specific data is typically split into non-overlapping support and query sets: ${\mathcal{D}^{(\mathrm{S})}_{\mathcal{T}}=\{(x_{i}^{(\mathrm{S})},y_{i}^{(
\mathrm{S})})\}_{i=1}^{N^{(\mathrm{S})}}}$ and
${\mathcal{D}^{(\mathrm{Q})}_{\mathcal{T}}=\{(x_{i}^{(\mathrm{Q})},y_{i}^{(
\mathrm{Q})})\}_{i=1}^{N^{(\mathrm{Q})}}}$ such that
$\mathcal{D}^{(\mathrm{S})}_{\mathcal{T}}\cap\mathcal{D}^{(\mathrm{Q})}_{
\mathcal{T}}=\emptyset$, with support and query set sizes, $N^{(\mathrm{S})}$ and $N^{(\mathrm{Q})}$, respectively.
The support set provides data for task-specific adaptation, and the query set provides data for evaluating the adapted performance.
The size of the support set is called the shot.
Note, even with an empty support set (*zero-shot*), the model can still adapt to the task using information from the query set,
i.e., related sequences without fitness scores.

Meta-learning for protein fitness prediction requires not just a single task, but multiple tasks, $\mathcal{D}={\mathcal{D}_{1},\dots\mathcal{D}_{\mathrm{T}}}$ over which to learn.
The full dataset of tasks, $\mathcal{D}$, can be seen as defining a distribution of tasks which can be split into training and test tasks in the usual way.
Concretely, for meta-learning in-context, the goal is to learn a function with parameters $\theta$ conditioned on the full support set and unlabelled inputs from query set,
$f_{\theta}(\{x_{i}^{(\mathrm{Q})}\}_{i=1}^{N^{(\mathrm{Q})}},\mathcal{D}^{(
\mathrm{S})}_{\mathcal{T}})$.
In our case, rather than directly predicting fitness values, we instead follow prior works that use a preference-based objective to rank the query set in order of fitness (Krause et al., ; Brookes et al., ; Hawkins-Hooker et al., ).

### 3.2 Metalic

Architecture Our work leverages the ProteinNPT architecture proposed by Notin et al. () as an in-context PLM for fitness prediction.
Here we briefly summarise the key elements, but defer the reader to Section A.2 and the original paper for full details.
The architecture, along with the data we use for meta-leaning, are illustrated in Fig. 2(a).

Protein sequences in both the support and query set are converted to per-residue embeddings using a pre-trained PLM; in our case we take the third layer of ESM2-8M (Lin et al., ).
Fitness scores in the query set are set to zero. All fitness scores are then projected to match the dimension of the residue embedding with a linear layer, summed with a learned embedding denoting if the row belongs to the support or query set, and finally concatenated along the sequence dimension.
Optionally, zero-shot fitness predictions from an auxiliary PLM can be embedded and concatenated in the same way, which we explore in Section 4.3.
This tensor is then processed via axial attention blocks (Ho et al., ), each of which applies self-attention separately along and across the sequences.
Axial attention reduces the computational complexity of self-attention from $O(K^{2}L^{2})$ to $O(K^{2}+L^{2})$, where $K$ is the shot and $L$ is the length of a protein.
Finally, a multi-layer perceptron conditions on the fitness embedding and mean-pooled sequence embedding to predict each query value, i.e. $\{v_{i}^{(\mathrm{Q})}\}_{i=1}^{N^{(\mathrm{Q})}}=f_{\theta}(\{x_{i}^{(\mathrm
{Q})}\}_{i=1}^{N^{(\mathrm{Q})}},\mathcal{D}^{(\mathrm{S})}_{\mathcal{T}})$, which is then used to rank the proteins by fitness.

Meta-Training Following prior works that use a preference-based objective (Krause et al., ; Brookes et al., ; Hawkins-Hooker et al., ),
we reframe the relative score prediction of two sequences as binary classification, predicting whether sequence $x_{i}^{(\mathrm{Q})}$ has a higher fitness than $x_{j}^{(\mathrm{Q})}$:

$$ $p\left(y_{i}^{(\mathrm{Q})}>y_{j}^{(\mathrm{Q})}\right)=\sigma\left(v_{i}^{( \mathrm{Q})}-v_{j}^{(\mathrm{Q})}\right),$ (1) $$

where $\sigma$ is a sigmoid function and the dependency of the query values on $\theta$ has been dropped for brevity. This classifier is optimized with respect to every pairwise comparison between sequences in the query set corresponding to the loss function:

$$ $\mathcal{L}(\theta,\mathcal{D}_{\mathcal{T}}^{(\mathrm{Q})},\mathcal{D}_{ \mathcal{T}}^{(\mathrm{S})})=-\sum_{i=1}^{N^{(\mathrm{Q})}}\sum_{\begin{ subarray}{c}j=1\\ j\neq i\end{subarray}}^{N^{(\mathrm{Q})}}\mathbb{I}\left(y_{i}^{(\mathrm{Q})}> y_{j}^{(\mathrm{Q})}\right)\log\sigma\left(v_{i}^{(\mathrm{Q})}-v_{j}^{( \mathrm{Q})}\right),$ (2) $$

where $\mathbb{I}$ is an indicator function.
Intuitively, this is optimising $N^{(\mathrm{Q})}\times(N^{(\mathrm{Q})}-1)$ binary classification problems.
Note that we only compute the loss over the query set to avoid encouraging memorization of the support set.
Adapting this to meta-learning (Fig. 2(a)), the objective becomes to find the parameterization that minimizes the loss across the task distribution,

$$ $\mathcal{J}(\theta,\mathcal{D})=-\mathbb{E}_{\mathcal{D}_{\mathcal{T}}\in \mathcal{D}}\mathbb{E}_{(\mathcal{D}_{\mathcal{T}}^{(S)},\mathcal{D}_{\mathcal {T}}^{(Q)})\in\mathcal{D}_{\mathcal{T}}}\mathcal{L}(\theta,\mathcal{D}_{ \mathcal{T}}^{(\mathrm{Q})},\mathcal{D}_{\mathcal{T}}^{(\mathrm{S})}).$ (3) $$

Fine-Tuning Metalic uses fine-tuning, during inference, in order to enable generalization, without having to account for fine-tuning during meta-training.
This process is depicted in Fig. 2(b).

The combination of in-context learning and fine-tuning creates a unique problem.
Since fine-tuning occurs at inference time, labels for the query set, $\{y_{i}^{(\mathrm{Q})}\}_{i=1}^{N^{(\mathrm{Q})}}$, are not available for training.
While labels for the support set, $\{y_{i}^{(\mathrm{S})}\}_{i=1}^{N^{(\mathrm{S})}}$, are available, propagating gradients from the support set would encourage memorization of the support set, since the labels are also passed as input in-context.
While prior methods that combine in-context and gradient-based meta-learning (Rusu et al., ; Vuorio et al., ) would encounter this issue, this problem is exacerbated for Metalic.
Whereas prior methods compress inputs to a representation with a constant number of dimensions, Metalic uses self-attention, which scales with the number of inputs, allowing them to be stored without compression.
Moreover, whereas prior methods use meta-gradients that could adjust the gradient update procedure so as to be useful for generalization and not memorization, Metalic does not take into account the fine-tuning process during meta-training.

We address the issue of memorization by sub-sampling from the support set.
The fine-tuning procedure is the same as during meta-training, with the exception that the support set is sub-sampled.
In order to compute updates on a single support set, the support set is sub-sampled into multiple smaller support and query sets, $\mathcal{D}^{(\mathrm{S^{\prime}})}_{\mathcal{T}}\subseteq\mathcal{D}^{(
\mathrm{S})}_{\mathcal{T}}$ and $\mathcal{D}^{(\mathrm{Q^{\prime}})}_{\mathcal{T}}\subseteq\mathcal{D}^{(
\mathrm{S})}_{\mathcal{T}}$.
Concretely, this corresponds to fine-tuning on unseen data using the objective:

$$ $\mathcal{J}(\theta,\mathcal{D}^{(\mathrm{S})}_{\mathcal{T}})=-\mathbb{E}_{( \mathcal{D}_{\mathcal{T}}^{(S^{\prime})},\mathcal{D}_{\mathcal{T}}^{(Q^{\prime })})\in\mathcal{D}^{(\mathrm{S})}_{\mathcal{T}}}\mathcal{L}(\theta,\mathcal{D} _{\mathcal{T}}^{(\mathrm{Q^{\prime}})},\mathcal{D}_{\mathcal{T}}^{(\mathrm{S^{ \prime}})}).$ (4) $$

After fine-tuning, Metalic conditions on the complete support set, allowing no data to go to waste.

Using Metalic’s unique combination of in-context meta-learning followed by fine-tuning, we enable the generalization of extensive fine-tuning, while also precluding expensive computation.
If a typical gradient-based meta-learning method requires $O(m)$ meta-gradients and $O(mn)$ regular gradients for meta-training, Metalic requires no meta-gradients and $O(m)$ regular gradients, constituting a linear reduction with superior performance to efficient alternatives, as demonstrated in Section 4.6.

## 4 Experiments

In this section, we evaluate Metalic on fitness prediction tasks from the ProteinGym benchmark (Notin et al., ).
We evaluate in the zero-shot setting with no support data, and the few-shot setting with limited support data.
To establish SOTA results in the zero-shot setting, we first compare to the predictions provided by ProteinGym for the baseline models.
To establish strong performance in the few-shot setting, since predictions are not provided, we train baselines from Hawkins-Hooker et al. ().
While Metalic does not achieve SOTA results in evaluations on proteins that have multiple mutations (multi-mutant proteins), we demonstrate that the performance grows as we add more meta-training tasks, providing a path forward for Metalic in the multi-mutant setting in the future.
We perform ablations of Metalic, to show the benefits of meta-learning, in-context learning, and fine-tuning.
Finally, we compare to the gradient-based method, Reptile (Nichol et al., ), to show that taking account of gradients during training is an unnecessary computational burden.

### 4.1 Experimental Setup

In our experiments, we focus on ProteinGym deep mutational scans.
Each task in ProteinGym each measures one property on a set of proteins that all differ by one amino acid, or multiple amino acids, from a reference wild-type protein.
We have 121 single-mutant tasks and 68 multi-mutant tasks from ProteinGym.
From these, we evaluate over eight held-out single-mutant tasks, and five held-out multi-mutant tasks, following Notin et al. (); Hawkins-Hooker et al. ().
This leaves 113 single-mutant and 68 multi-mutant tasks for training when evaluating single mutants (Sections 4.2 and 4.3), and 121 single-mutant and 63 multi-mutant tasks for training when evaluating multiple mutants (Section 4.4).
All fitness values are standardized by subtracting the mean and dividing by the standard deviation by task.
Note that there are additional tasks in ProteinGym we do not consider.
Specifically, we do not consider multi-mutant tasks that have overlapping proteins with single-mutant tasks, to make evaluation more difficult.
We also ignore additional tasks in which the maximum protein length is $>$ 750, to fit the backward pass on an Nvidia A100-80Gb device.

We use a query set size of $N^{(\mathrm{Q})}=100$, and the size of the support set is determined by our evaluation setting and is one of three sizes: $N^{(\mathrm{S})}=$ 0, 16, or 128.
We also use an additional set of 128 points just for early stopping of the fine-tuning process, for all models, following the implementation of Hawkins-Hooker et al. ().
We then evaluate remaining points in the task, with a maximum of 2,000 points total, by dividing the data into multiple query sets.
If the model can fit a larger query size, as in non-meta-learning baselines, then we pass the remaining data as a single query set.
If the data is not divisible by the query set size, it is left out from evaluation.
However, we sample the support data over three independent samples, avoiding systemic exclusion.

All evaluation uses the Spearman rank correlation, in line with prior work (Notin et al., ; Truong Jr & Bepler, ; Hawkins-Hooker et al., ).
We compute the Spearman correlation per task, and then average over tasks.
For all evaluations of our models, we compare over three seeds for training and report the mean and standard deviation.
Each context, consisting of a support and query set, consists of $\leq$ 171,000 tokens.
We meta-train for 50,000 updates and fine-tune for 100.
Fine-tuning uses the same procedure, including an Adam optimizer and cosine learning rate scheduler, but fine-tuning skips the learning rate warm-up, so the scheduling has little effect.
Using a single Nvidia A100-80Gb, training our model takes roughly 2 to 8 days per seed, depending on support size and frequency of fine-tuned evaluation.

### 4.2 Zero-Shot

The first setting we evaluate is the zero-shot performance of our model, with no support set for fine-tuning.
In Table 1 we compare against predictions provided by ProteinGym for each baseline to compute the zero-shot Spearman correlation ($\rho$).
We compare to provided predictions, on our data splits, to enable a fair comparison to the strongest models available without retraining each baseline from scratch ourselves.
We include the best performing model, and notable models, as baselines.
We find that Metalic outperforms every reported baseline and is SOTA at zero-shot prediction.

**Table 1: Spearman correlation in the zero-shot setting. Results are computed using predictions provided by ProteinGym. Using a single PLM, in the zero-shot setting, renders the baselines deterministic. For comparison, we report our best model, and the mean and standard deviation, for quantifying the variation in meta-training. Metalic achieves SOTA performance in either case.**
| Model Name | Spearman Correlation |
| --- | --- |
| Metalic (max) | .484 |
| Metalic (mean) | .482 $\pm$ .002 |
| VESPA | .464 |
| TranceptEVE-Medium | .457 |
| ESM1-v-650M | .437 |
| Tranception-Medium | .427 |
| Progen2-Medium | .419 |
| ESM2-650M | .399 |
| MSA Transformer | .398 |
| ESM-IF1 | .365 |
| ESM2-8M | .121 |

Our method significantly outperforms strong baselines with many more parameters, such as ESM1-v-650M.
The 8 million parameter PLM used by our method, ESM2-8M, without Metalic, achieves a score of only $.121~{}\rho$, demonstrating the large contribution of our meta-learning procedure.
The next strongest method after ours is VESPA (Marquet et al., ).
VESPA optimizes PLM embeddings to predict a distinct set of binary annotations from 9,594 proteins, and uses these features to predict protein fitness by comparing to a reference wild-type embedding.
Note, unlike our method, VESPA relies on strong assumptions to generalize and is specific to the zero-shot setting.

The strong performance of our method in the zero-shot setting can be attributed to meta-learning.
Since there is no data for fine-tuning, the zero-shot performance increase over ESM2-8M derives entirely from our meta-training procedure.
In this case, other methods generally assume that the PLM likelihood of a mutation correlates with fitness, whereas our model learns to make use of the information contained in PLM embeddings to make predictions zero-shot.
Moreover, our method still conditions on an unlabeled query set, and the protein embeddings in that query set, which allow for meta-learning a form of in-context unsupervised adaptation.

### 4.3 Fine-Tuning Results

In Table 2 we report Spearman correlation with a support set of size $N^{(\mathrm{S})}=$ 0, 16, and 128, and we compare to baselines that we train and evaluate ourselves over three random seeds.
We re-train these methods using the models, following Hawkins-Hooker et al. (), to provide a comparison over multiple seeds between these methods in a range of practical low-data settings.
All models use the same preference-based loss function as Metalic, for a fair comparison, and none use ensembling.

We also train Metalic-AuxIF, which allows Metalic to condition on zero-shot predictions from an additional PLM, as introduced in Section 3.2.
We choose ESM-IF1 (Hsu et al., ) as the auxiliary input, since it contains embeddings derived from inverse folding that summarize protein structure.
Note that Metalic-IF1 uses 170 million parameters, whereas Metalic uses 36 million (including the ESM2-8M embedding), so the auxiliary predictions significantly increases the total parameters.

Again, we find that Metalic has the strongest performance in the 0-shot and 16-shot settings, and has comparably strong performance in the 128-shot setting, with 18 times fewer parameters.
Moreover, Metalic also outperforms contemporary models, such as PoET (Truong Jr & Bepler, ), that make use of additional evolutionary information, in the form of multi-sequence alignment, and in-context learning.
Note that ESM2-8M and ProteinNPT are the worst performing methods in the few-shot setting.
This results suggests that the effectiveness of Metalic does not come from the ESM2-8M embedding, nor the ProteinNPT architecture, but rather from the meta-learning itself.

**Table 2: Spearman correlation for the 0, 16, and 128-shot setting. Baseline results are re-computed. Standard deviation is provided over three seeds. Metalic matches or exceeds all baselines.**
| Model Name | $n=0$ | $n=16$ | $n=128$ |
| --- | --- | --- | --- |
| Metalic | .482 $\pm$ .002 | .484 $\pm$ .001 | .552 $\pm$ .009 |
| Metalic-AuxIF | .498 $\pm$ .008 | .500 $\pm$ .002 | .556 $\pm$ .005 |
| ESM1-v-650M | .384 $\pm$ .000 | .452 $\pm$ .000 | .553 $\pm$ .000 |
| ESM2-8M | .105 $\pm$ .000 | .226 $\pm$ .000 | .406 $\pm$ .000 |
| PoET | .416 $\pm$ .003 | .475 $\pm$ .026 | .588 $\pm$ .006 |
| ProteinNPT | N/A | .192 $\pm$ .003 | .443 $\pm$ .003 |

Additionally, we see that Metalic-AuxIF, improves results.
This result is significant because ESM2-8M and ESM-IF1, the two PLMs used by Metalic-AuxIF, are the worst performing methods in the zero-shot setting (Table 1), with ESM2-8M also weak in the few-shot setting.
Thus, the effectiveness of Metalic-AuxIF comes entirely from meta-learning: We can learn how and when to rely on features from each of these weak predictors of protein fitness, to combine them into a strong predictor.

Consistent with the motivation of meta-learning, results are strongest when the data is most limited.
Meta-learning adds an additional training stage to learn prior beliefs and inductive biases from related data.
The more limited, the more relying on prior data is useful.

### 4.4 Multiple Mutants

Figure: (a) Table 3a: Multi-mutant Spearman correlation in the zero-shot setting. Results are computed using predictions provided by ProteinGym on tasks with multiple mutations. Metalic is competitive, but not better than all baselines, due to lack of sufficient datasets with multiple mutants.

In this section we evaluate Metalic on tasks where proteins have multiple mutations.
In Table 3(a), we see that Metalic has strong performance, but does not outperform all baselines.
However, with auxiliary inverse folding predictions (Metalic-AuxIF), Metalic is only outperformed by ESM-IF1 itself.
These results indicate that ESM-IF1 is particularly strong on the eight held-out evaluation tasks, and that we can recover most of its performance by incorporating its predictions into our model.
Note that while ESM-IF1 is strong in the multi-mutant setting, it is among the worst performing models in the single-mutant setting.
In contrast, our method is strong in both settings, since it can learn to leverage ESM-IF1 predictions when they are helpful, and ignore them when they are not.

By comparing across all shot settings, we can see that Metalic is competitive in the multi-mutant setting (Fig. 3(b)), but no longer SOTA, which we hypothesize is due to limited multi-mutant training data.
In Table 4(a) and Fig. 4(b), we explore this hypothesis.
Specifically, we see the performance of Metalic increases as the amount of training data, measured in tasks, increases.
This trend holds true even when adding single mutant tasks, which are significantly different from the testing data.
Providing more data is a path forward for strengthening the multi-mutant results.

Figure: (a) Table 4a: Multi-mutant Spearman correlation by training data. Results are computed on eight tasks with multiple mutations with different amounts of single- and multi-mutant training tasks. We see performance of Metalic increases with more data, giving a path to improve results with future data collection. This trends holds even when we add single-mutant data, which is significantly different from the multi-mutant test data.

### 4.5 Ablations

Here we report ablations of Metalic, evaluating on single-mutants (Table 5).
Spearman correlation is reported in the zero-shot and 128-shot settings.
The results justify all components of our method.

Most detrimental to performance is removing meta-training from Metalic (NoMetaTrain).
This ablation is identical to Metalic, with the exception that there is no additional meta-learning stage over multiple protein landscapes.
We see that without meta-learning, the initial zero-shot predictions have near zero correlation with the fitness and that the 128-shot predictions are critically impaired.

**Table 5: Ablations in the 0 and 128-shot setting. Results show the importance of fine-tuning, in-context learning, meta-training, and additional training tasks as an augmentation.**
| Model Name | $n=0$ | $n=128$ |
| --- | --- | --- |
| Metalic | .482 $\pm$ .002 | .552 $\pm$ .009 |
| Metalic-NoFT | .482 $\pm$ .002 | .488 $\pm$ .012 |
| Metalic-NoPref | .465 $\pm$ .011 | .520 $\pm$ .006 |
| Metalic-NoAug | .464 $\pm$ .011 | .558 $\pm$ .002 |
| Metalic-NoICL | .441 $\pm$ .002 | .529 $\pm$ .002 |
| Metalic-NoMetaTrain | -.046 $\pm$ .018 | .346 $\pm$ .003 |

We also ablate the ability to attend to the rest of the proteins in context by turning off the column attention in the axial attention layers (NoICL), we ablate the fine-tuning stage of training (NoFT), we ablate the preference-based loss by using mean squared error instead (NoPref), and we ablate the augmentation of single-mutant training data with multi-mutant data (NoAug).
Note that when we remove the fine-tuning, and have a non-zero support size, we allow the early stopping data to be passed in-context, to not unfairly advantage fine-tuning with additional data.
Ablating in-context learning decreases performance in both settings, indicating an ability to adapt in an unsupervised fashion even from the query set alone.
Ablating fine-tuning decreased performance in the 128-shot setting, indicating the need for fine-tuning for generalization to held-out data.
Ablating the preference-based loss decreased performance in both settings, in line with recent literature (Krause et al., ; Brookes et al., ; Hawkins-Hooker et al., ).
Finally, ablating the multi-mutant augmentation decreased performance in the zero-shot setting, with comparable 128-shot performance, indicating the benefit of additional meta-training data when fine-tuning data is most limited.

### 4.6 Gradient-Based Meta-Learning

In this section, we compare Metalic to the same architecture but trained with Reptile (Nichol et al., ), an efficient method for gradient-based meta-learning.
Unlike Metalic, Reptile does not use in-context learning and modifies the outer-loop during meta-training to take into account the subsequent fine-tuning.
While accounting for the fine-tuning during meta-training comes with increased computational costs and can increase bias and variance (Vuorio et al., ), Reptile provides a simplified algorithm that can be run more efficiently.
The primary differences between Reptile and Metalic are that Reptile adjusts the meta-learning loss to account for fine-tuning during meta-training, while Metalic performs in-context learning.
Reptile details are provided in Section A.1.

**Table 6: Comparison to Reptile. Reptile (Nichol et al., ) is a gradient-based meta-learning method. Results are reported after 15,000 steps, with Metalic additionally reported after 45,000 steps, to allow for an equal number of gradient computations as the methods using Reptile. Results show that accounting for fine-tuning during meta-training using Reptile is unnecessary.**
| Model Name | Meta-Training Steps | Total Gradient Computations | $n=128$ |
| --- | --- | --- | --- |
| Metalic (45k) | 45,000 | 45,000 | .553 $\pm$ .003 |
| Metalic (15k) | 15,000 | 15,000 | .512 $\pm$ .003 |
| Reptile-3-3 | 15,000 | 45,000 | .504 $\pm$ .009 |
| Reptile-3-100 | 15,000 | 45,000 | .528 $\pm$ .004 |
| Metalic-Reptile | 15,000 | 45,000 | .532 $\pm$ .006 |

Although Reptile is more compute-efficient than other gradient-based methods, it still performs gradient updates in the inner loop during meta-training, making it computationally expensive relative to Metalic.
Our method uses 100 updates on the support data for fine-tuning.
Reptile uses inner-loop gradient updates during both meta-training and fine-tuning.
Due to compute limitations, we cannot use 100 steps for each forward pass of meta-training.
Thus, we evaluate Reptile with 3 inner-loop gradient steps during meta-training and 3 during fine-tuning, so that train and test match (Reptile-3-3), and we evaluate Reptile with 3 inner-loop gradient steps during meta-training and 100 during fine-tuning, so that test time matches our method (Reptile-3-100).
Note that even Reptile-3-3 uses three times more compute than Metalic.
Finally, we evaluate Reptile-3-100 with Metalic, to see whether they can be used in conjunction (Metalic-Reptile).
We train all methods for 15,000 steps, given the increased computation, and also report Metalic after 45,000 steps, to match the same number of gradient computations as the Reptile models.

Results are reported in Table 6.
We evaluate in the 128-shot setting.
Note that Reptile is not applicable in the zero-shot setting, as it requires some data for fine-tuning.
We observe that without controlling for the amount of computation, Metalic (15k) performs similarly to Reptile, with a Spearman correlation between Reptile-3-3 and Reptile-3-100.
We also observe that Metalic, in conjunction with Reptile (Metalic-Reptile), outperforms Reptile and Metalic independently.
However, the improvement over Metalic is marginal, and the increased computation (by a factor of 3) is large.
Controlling for the total number of gradient computations,
Metalic (45k) achieves a Spearman correlation much higher than the combined Metalic-Reptile.
Thus, while Reptile can be used in conjunction with Metalic, it may not be worth the added computation.
This result is in line with previous work showing that accounting for gradient updates during meta-training (in their case, without in-context learning) can be detrimental when there are a limited number of tasks for meta-training or a limited amount of data for the inner-loop (Gao & Sener, ; Triantafillou et al., ).
Collectively, this evidence further justifies fine-tuning only after meta-training.

## 5 Analysis

Figure: (a) Untrained Attention
Refer to caption: extracted/5917705/figs/col_attention_heatmap_1.png

Here, we briefly investigate the in-context learning abilities of Metalic via attention maps.
In Section 4, we show that in-context learning is vital to Metalic by ablating the attention between proteins and showing decreased performance.
Notably, the in-context learning was beneficial not only in the few-shot setting, but also in the zero-shot setting.
This suggests an interesting phenomenon: The emergence of unsupervised in-context learning from the query set alone.
To confirm this phenomenon, in Fig. 5 we show the attention maps in the axial attention layers between proteins in the query set.
Over the course of training, we observe the emergence of bright vertical lines, indicating some significantly informative proteins to which all others proteins attend.
Moreover, we observe no rows that are entirely dark off the diagonal entries.
Thus, no protein attends only to itself.
Together, these results further corroborate the necessity of in-context meta-learning.

## 6 Conclusion

In this paper, we have demonstrated state-of-the-art results on a standard protein fitness prediction benchmark in low-data settings.
In order to do so, we proposed Metalic, which makes use of both in-context meta-learning and subsequent fine-tuning.
Critically, we have demonstrated the ability of meta-learning to take advantage of additional data from other proteins fitness prediction tasks, while remaining computationally tractable by deferring fine-tuning to test time alone.
Unique within the meta-learning literature, we show that in-context meta-learning provides a useful initialization for further fine-tuning, and can make use of test time data for both fine-tuning and in-context learning.
Metalic additionally demonstrates the ability to learn from the query set alone (zero-shot), performing unsupervised in-context learning.
Future work could investigate leveraging Metalic’s unique in-context learning ability to act as an auto-regressive fitness model (i.e., a world model in the meta-reinforcement learning setting (Beck et al., )), for optimally trading off exploration and exploitation when designing novel proteins.
Given its efficacy at leveraging additional data, we believe that meta-learning in-context will play a crucial role in advancing protein fitness prediction, with Metalic being a foundational step in that direction.

## Appendix A Appendix

### A.1 Reptile Details

This section provides additional details on how Reptile (Nichol et al., ) works. In order to account for fine-tuning during meta-training, gradient-based methods generally compute a meta-gradient that requires the computation of higher order derivatives, which can be computationally intractable for a large model.
The costs can be especially burdensome when many gradient steps are needed for out of distribution adaptation.
For this reason, Reptile avoids meta-gradients by changing optimization during meta-training.
Here, the new parameters are updated, not by gradient descent, but rather by moving toward the mean, after each theta is adapted to a task, $\theta_{\mathcal{T}}$, by fine-tuning on that task.
From time-step $t$ to $t+1$ of this outer-loop optimization process, Reptile can be written:

$$ $\theta^{t+1}=\theta^{t}+\beta\mathbb{E}_{\mathcal{D}_{\mathcal{T}}\in\mathcal{ D}}[\theta_{\mathcal{T}}^{t}-\theta^{t}].$ (5) $$

We had to choose several implementation details for Reptile.
First, note that Reptile sub-samples the support set during fine-tuning and has no distinct query set.
In our implementation, we sub-sample mini-batches of size 50 to match the query size of Metalic for sub-sampling in the 128-shot setting.
Reptile also has several options for the outer loop optimization.
We choose to use the batched version with a batch size of 4 to match Metalic.
Rather than using the direct update given in Equation 5, we take the difference over the learning rate, $(\theta_{\mathcal{T}}^{t}-\theta^{t})/\alpha$ as an approximation of a gradient to be used with the Adam optimizer, as suggested by Nichol et al. () and validated in Section A.2.

### A.2 Hyper-Parameters

Since we build upon the axial attention of ProteinNPT (Notin et al., ), we follow their choice for most hyper-parameters, with a few exceptions.
Most notably, in our experiments, we use the third layer of ESM2 as an embedding for each protein, given the strong performance and reduced number of parameters.
Additionally, we use a ranking loss from Hawkins-Hooker et al. (), and do not use the CNN or additional inputs (such as zero-shot predictions) from ProteinNPT.
We likewise found conditioning on the wild-type unhelpful.
The same set of hyper-parameters are used for each setting: 0-shot, 16-shot, and 128-shot, and for multi-mutant results.
We used the same learning rate for fine-tuning Metalic as for meta-training.
We found the default ProteinNPT learning rate to be too large, and decreasing by a factor of 5 to be sufficient.
Complete details on hyper-parameters used are available in Table 7.
We tuned relatively few of the hyper-parameters of our method, and mostly tuned over a single seed.
There is likely room for improvement in the hyper-parameter selection of Metalic.

**Table 7: Hyper-Parameters for Metalic. Where different, an earlier set of hyper-parameters used for ablations and Reptile comparisons are given in parentheses.**
| Hyper-Parameter | Description | Value |
| --- | --- | --- |
| Training Steps | The total number of training steps in meta-training | 50,000 |
| Warm-Up Steps | The number of training steps spent linearly warming up, preceding cosine decay | 5,000 |
| Batch Size | The number of contexts evaluated per training step. Note that gradient accumulation is used for each context in the batch, so this scales training time linearly. | 4 |
| Weight Decay | Weight decay applied to non-bias parameters only | 5e-3 |
| Learning Rate | The learning rate for meta-training and fine-tuning | 6e-5 |
| Min LR Fraction | The minimum fraction of the LR maintained during the cosine decay in learning rate scheduling | 1e-5 |
| Adam Eps | The epsilon value for the Adam optimizer | 1e-8 |
| Adam Beta1 | The beta1 value for the Adam optimizer | 0.9 |
| Adam Beta2 | The beta2 value for the Adam optimizer | 0.999 |
| Gradient Clip Value | The maximum norm allowed for the gradient | 1.0 |
| ESM embed model | The full name for the ESM2 model used | esm_t6_8M_UR50D |
| ESM embed layer | The layer from the ESM2 model used as an embedding | 3 |
| Number Fine-tune Steps | The number of gradient updates for fine-tuning after meta-training | 100 |
| Num ProteinNPT Layers | The number of layers using axial attention, as in ProteinNPT | 5 |
| Condition on Pooled Sequence | Whether each sequence is pooled or ignored after axial attention | True |
| MLP Layer Sizes | The number and size of fully connected layers after axial attention | [768,768,768,768] |
| Embed Dim | The embedding dimension for all inputs including the protein sequences and fitness values | 768 |
| Axial Forward Embed Dim | The hidden size of the feed-forward layer within the ProteinNPT layer | 400 |
| Attention Heads | The number of heads in self-attention | 4 |
| Dropout Prob | The probability of dropout during training and fine-tuning for layers other than axial attention | 0.0 |
| Attention Dropout | The probability of dropout during training and fine-tuning for axial attention layers | 0.1 |
| Num Single Tasks | The total number of single-mutant tasks available. These tasks are included for meta-training even when testing on multi-mutants. Eight are held-out for evaluation when evaluating single-mutant performance. | 121 |
| Num Multiple Tasks | The total number of multi-mutant tasks available. These tasks are included for meta-training even when testing on single-mutants. Five are held-out for evaluation when evaluating multi-mutant performance. | 68 |
| Warm-up During Fine-tuning | Whether or not to use the linear warm-up from the learning rate scheduler during fine-tuning. | False |

For the majority of baselines no tuning was required, other than Reptile.
For the baselines, we used reference predictions for Table 1 and reference implementations for Table 2, neither of which required tuning.
For Reptile, the update in the outer-loop can be uses as written in Equation 5, or $(\theta_{\mathcal{T}}^{t}-\theta^{t})/\alpha$ can be interpreted as a gradient for use with the Adam optimizer Nichol et al. ().
We experiment with both and find the use of the Adam optimizer to be superior in performance.
All results in the main body use the Adam optimizer for Reptile.
In order to evaluate the method without the Adam optimizer, we re-tune $\beta$ in Equation 5, which is the outer-loop learning rate.
We leave the inner-loop learning rate, $\alpha$, fixed at the learning rate for Metalic, $3e^{-5}$, which uses the same learning rate for the inner- and outer-loop.
Given the increased computational cost of Reptile, we use a single seed over three learning rates for 10,000 steps.
We tune over the following learning rates ($\beta$): the learning rate for Metalic, which is $3e^{-5}$; a learning rate of 1, which corresponds to the learning rate of Metalic if you interpret $(\theta_{\mathcal{T}}^{t}-\theta^{t})/\alpha$ as the gradient in a gradient-based update (Nichol et al., ); and a learning rate of $\frac{1}{n}$, where $n$ is the number of inner-loop updates (3 in this case), which corresponds to the learning rate of Metalic if you interpret $(\theta_{\mathcal{T}}^{t}-\theta^{t})/(\alpha n)$ as the gradient in a gradient-based update(^1^11he standard gradient update, $\theta^{t+1}=\theta^{t}+\alpha\nabla$, with $\nabla=(\theta_{\mathcal{T}}^{t}-\theta^{t})/(\alpha n)$, gives $\theta^{t+1}=\theta^{t}+\alpha(\theta_{\mathcal{T}}^{t}-\theta^{t})/(\alpha n)$, which is the same update as Equation 5 with $\beta=\frac{1}{n}$, omitting the expectation for brevity.).
We found a learning rate of 1 to perform best, in line with the outer loop learning rate of Metalic and the interpretation of the gradient from Nichol et al. ().
For Adam, there is no $\beta$, and we do not re-tune Adam’s outer-loop learning rate.
When using Adam, the gradient is interpreted as $(\theta_{\mathcal{T}}^{t}-\theta^{t})/\alpha$.
The results of tuning the learning rate without Adam, suggests a learning rate of $3e^{-5}$ with this gradient interpretation.
Indeed, we confirm that when using Adam with this learning rate, and this gradient interpretation, it works better than using the alternative Reptile update rule with any learning rate.
Results of the optimizer and learning rate tuning are presented in table Table 8.
Note that these results used an earlier set of hyper-parameters.

**Table 8: Reptile tuning results for the 128-shot setting. Results are trained over 10,000 steps for one seed each. The Adam optimizer performs the best. Both Adam, with an outer-loop learning rate of $3e^{-5}$, and Equation 5, with $\beta=1$, correspond to the same scale in the outer loop.**
| Model Name | Adam ($3e^{-5}$) | $\beta=1.$ | $\beta=.333$ | $\beta=3e^{-5}$ |
| --- | --- | --- | --- | --- |
| Reptile-3-3 | .452 | .420 | .350 | .091 |
| Reptile-3-100 | .472 | .444 | .396 | .181 |
| Metalic-Reptile | .483 | .476 | .411 | .190 |

### A.3 Fine-Tuning Warm-Up

Metalic does not use a warm-up period for fine-tuning as it does for meta-training.
This decision was made since the warm-up normally occurs for 5,000 steps, but fine-tuning only occurs for 100 steps.
As a compromise, we evaluate a warm-up period for fine-tuning, after taking a single gradient step with the full learning rate (Metalic-FTWarmUp).
We still see that not using a warm-up for fine-tuning is superior.
Results can be see in Table 9.

**Table 9: Spearman correlation with and without warm-up.**
| Model Name | $n=128$ |
| --- | --- |
| Metalic | .552 $\pm$ .009 |
| Metalic-FTWarmUp | .539 $\pm$ .007 |

### A.4 More Multi-Mutant Results

In this section we present a table of additional results for tasks with multiple mutants, including the Metalic-AuxIF models.

In Table 11, we see the performance of Metalic increases as the amount of training data, measured in tasks, increases, providing more data is a path forward for strengthening the multi-mutant results.
Here, we also see that the trend changes for Metalic-AuxIF, since it performs better with more data, but only if that data is multi-mutant data.
This makes sense, since the auxiliary ESM-If1 predictions are most effective for multi-mutants, and training on single-mutants could discourage relying on them.

**Table 10: Spearman correlation for the 0, 16, and 128-shot setting for multi-mutant tasks.**
| Model Name | $n=0$ | $n=16$ | $n=128$ |
| --- | --- | --- | --- |
| Metalic (Both) | .436 $\pm$ .011 | .484 $\pm$ .001 | .670 $\pm$ .003 |
| Metalic-AuxIF (Both) | .480 $\pm$ .007 | .500 $\pm$ .002 | .693 $\pm$ .006 |
| Metalic-AuxIF (MultiTrain) | .533 $\pm$ .012 | .577 $\pm$ .008 | .692 $\pm$ .003 |
| ESM1-v-650M | .426 $\pm$ .000 | .557 $\pm$ .000 | .644 $\pm$ .000 |
| ESM2-8M | .368 $\pm$ .000 | .484 $\pm$ .000 | .596 $\pm$ .000 |
| PoET | .588 $\pm$ .014 | .638 $\pm$ .001 | .734 $\pm$ .006 |
| ProteinNPT | N/A | .367 $\pm$ .003 | .624 $\pm$ .003 |

**Table 11: Multi-mutant Spearman correlation by training data. Results are computed using predictions provided by ProteinGym on five tasks with multiple mutations with different amounts of single- and multi-mutant training tasks. We see performance of Metalic increases with more data, giving a path for future data collection. We also see here that Metalic-AuxIF performs better when trained with only multi-mutant data. This makes sense, since the augmented predictions are most effective for multi-mutants, and training on single-mutants could discourage relying on them. An asterisk (^∗) represents the default training data regime for that method.**
|  |  |  | Multiples |
| --- | --- | --- | --- |
| Model Name | Single-Mutant Tasks | Multi-Mutant Tasks | $n=0$ |
| Metalic (Both)^∗ | 121 | 63 | .436 $\pm$ .011 |
| Metalic (MultiTrain) | 0 | 63 | .376 $\pm$ .018 |
| Metalic (MultiTrainHalf) | 0 | 30 | .307 $\pm$ .022 |
| Metalic-AuxIF (Both) | 121 | 63 | .480 $\pm$ .007 |
| Metalic-AuxIF (MultiTrain)^∗ | 0 | 63 | .533 $\pm$ .012 |
| Metalic-AuxIF (MultiTrainHalf) | 0 | 30 | .517 $\pm$ .021 |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: * How does performance change when training and evaluating on a separate set of proteins?
* How do inference and training costs of Metalic compare to the Reptile ablations?
* To clarify the architecture, would it be possible to have a Vaswani et al. 2017 style figure in the appendix of the exact computations that occur?

Overall the work is quite strong, and would benefit from a simple experiment to further prove its robustness. 

Vaswani et al. 2017, "Attention Is All You Need"
soundness: 3
strengths: Metalic achieves success in both providing a compute efficient alternative to meta-learning and also provides strong performance on downstream tasks. It’s also noteworthy the parameter efficiency that Metalic achieves in the process.

The work is also highly reproducible. Each experiment could be reproduced by a reader from the text alone (barring which datasets were and were not used). In the appendix, readers can find a comprehensive set of hyperparameters and training details used in each experiment.

When comparing metalic, it uses relevant choices in baseline models and also analyzes where and how performance degrades (e.g. the case of multi-mutants). In addition, the ablations show the utility of the different components of the method, even displaying its ability to outperform traditional meta-learning methods like Reptile. The experiments are very thorough in justifying all choices made during modeling.
summary: This work introduces Metalic, a framework for performing protein sequence fitness prediction in both in-context and few-shot regimes. It proposes to use in-context-learning as a replacement for gradient updates used in the traditional inner loop of meta-learning. While using the smallest ESM-2 model, it achieves strong performance on the chosen fitness tasks.
weaknesses: The only major weakness that appears in this work is the limited number of datasets used in evaluating the proposed methods. ProteinGym performance fluctuates a lot amongst tasks, hence it would be nice to decouple the comparatively small number of datasets, 13 in total, for evaluation from performance. A simple experiment selecting different splits and training replicates in a similar manner to a partial k-fold cross validation would be convincing.

On a more minor note, the exact implementation of the algorithm would benefit from a slightly more detailed treatment. Though it can be inferred from Figure 2a, it would benefit from some support either pictorially or mathematically in the exact computations that must be performed.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: * Are there some tasks that are easier to learn vs. harder to learn? It would be interesting to see a breakdown analysis.
* Do authors have hypotheses for why the performance gap between Metalic and baselines is so much more pronounced for zero-shot vs. few shot tasks?
soundness: 3
strengths: * The idea is relatively under-explored in protein ML despite ICL having shown strong performance in other fields of ML, and the importance of few-shot learning.
* Paper is fairly clearly written
* Attention map analysis is a great addition to the work! Would love to see follow up on this.
summary: This paper introduces Metallic, which performs few-shot (N=16 or N=128) or zero-shot fitness prediction using in-context meta-learning, using the ProteinGym benchmark. The model is meta-trained across many protein prediction tasks, to enable in-context learning, then finetuned on the support set, where the support set defines the number of "shots" for the few-shot task. Evaluations uses pairwise ranks (i.e. if model correctly prefers one mutant to the other). The meta-training procedure yields models that fare better in the finetuned few-shot setting than baselines (ESM1, etc.)
weaknesses: * In Table 1, 2, and 3a, only ESM2-8M is used for comparison. What happens if we compare to the larger ESM2 models? Can we get better zero-shot or few-shot performance via scaling pretraining parameters?
* Results in Figure 3b and 3a don't seem too much better than baselines.
* More rigorous examination of the effects of varying query set size, etc. might be quite insightful
* Some papers have pointed out that protein likelihoods don't always indicate protein fitness, and might be due to the training data biases [1,2].

Minor / meta comment:
I think that though thinking about ICL for proteins is a good idea, and authors demonstrate thoughtfulness in how this should be done. However, there might be more interesting tasks to examine than likelihood & protein fitness. There has also been working showing that "evotuning" (finetuning on evolutionary neighbors) can improve performance [2,3]. In a real-world setting, if I have a protein for which my pretrained model is underperforming on, I'm not sure if I would prefer to use a meta-trained model or a evo-tuned model.

From a writing perspective, authors might be able to increase impact by outlining real-world drug discovery scenarios where this is relevant.

[1] Protein language models are biased by unequal sequence
sampling across the tree of life. Ding et al., 2024
[2] Protein Language Model Fitness Is a Matter of Preference. Gordon et al., 2024
[3] Unified rational protein engineering with sequence-based deep representation learning. Alley et al., 2024

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1-The authors assume that the correlation between PLM likelihood and protein fitness may not be reliable. How does Metalic perform when the assumed correlation is weaker? Would additional analysis on this correlation help refine the model's approach?

2-In-context learning relies heavily on the distribution of tasks used during meta-training. How does Metalic handle tasks that are significantly different from the meta-training distribution, and are there strategies that could improve its generalization to such out-of-distribution tasks?

3-Given that Metalic does not account for fine-tuning during meta-training, how does this omission affect the model's ability to adapt to tasks with extensive fine-tuning requirements? Would incorporating some level of fine-tuning awareness during meta-training improve results?
soundness: 3
strengths: 1-By integrating in-context meta-learning with subsequent fine-tuning, Metalic can generalize well to new tasks, even though fine-tuning is not accounted for during meta-training. Moreover, the model is more computationally efficient than state-of-the-art approaches.

2-Metalic achieves strong performance in zero-shot and few-shot protein fitness prediction tasks, which is particularly valuable given the limited availability of labeled data for many protein-related tasks.
summary: The paper titled introduces Metalic, a meta-learning approach for predicting protein fitness, which is a measure of protein properties like stability and binding affinity. Predicting these properties accurately is challenging due to the scarcity of labeled data. Metalic improves on this by applying meta-learning to a distribution of related protein fitness tasks, enabling positive transfer to unseen tasks. The approach combines in-context meta-learning and fine-tuning without accounting for fine-tuning during meta-training, which helps to generalize well even with limited data. Experiments show that Metalic achieves state-of-the-art performance in zero-shot and few-shot settings on the ProteinGym benchmark, using 18 times fewer parameters than current models.
weaknesses: 1-I anticipate a more in-depth analysis of the correlation between PLM probability and fitness values in the context of meta-learning, especially when the authors assume that this correlation may not be reliable.

2-Although in-context learning is a strength, its effectiveness diminishes for tasks that deviate significantly from the distribution of tasks used during meta-training, potentially limiting generalization to out-of-distribution tasks. Experimental results on some OOD tasks can help readers better understand the paper.

3-The paper primarily compares Metalic to a few gradient-based meta-learning methods (e.g., Reptile) but does not extensively benchmark against other recent meta-learning approaches that incorporate different strategies for handling fine-tuning or leveraging unsupervised data.

4-Metalic shows strong performance on single-mutant tasks but falls short of the state-of-the-art on multi-mutant tasks. The authors attribute this to the limited availability of multi-mutant training data, emphasizing the need for more diverse datasets to fully realize the method's potential. A more in-depth analysis of this limitation would be valuable.

5-While the ProteinGym benchmark is a standard for protein fitness prediction, additional evaluations on other benchmarks would provide a more comprehensive understanding of Metalic's strengths and weaknesses across a broader range of protein-related tasks.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 4
questions: While the paper is fairly clearly written, there are a few pieces of information I haven't been able to find:

1. Why did you stop at 8M? How much more computationally intensive would it be to run experiments with the 35M model?
2. How many tasks are presented in-context to generate the figures in e.g. Table 1? What data is used in-context at evaluation time? Is there overlap with the data used during the meta-learning phase?
3. "Moreover, our method still conditions on an unlabeled query set, and the protein embeddings in that query set, which allow for meta-learning a form of in-context unsupervised adaptation." -> Is there any evidence in the paper that this occurs? Would be cool to see an ablation of the data from other tasks to isolate this.
soundness: 4
strengths: The paper is clearly written, includes a range of baselines and ablations, and also attacks an important problem. In some sense, protein fitness prediction is an excellent application for meta-learning, in that it's one of the few settings I can think of where there is an almost arbitrarily broad distribution of tasks that are nevertheless closely related to one another.
summary: The authors present Metalic, a new approach to predicting properties of protein sequences. Metalic involves two phases: a "meta-learning" phase, during which a PLM backbone is fine-tuned to learn from data from multiple related fitness prediction tasks presented in-context, and then a per-task "fine-tuning" phase, where models are fine-tuned at inference time on a small amount of data specific to the task at hand. The authors show that Metalic outperforms baselines in the zero-shot setting for single mutations and is competitive for multiple mutations.
weaknesses: I'm a bit conflicted about this paper. On the one hand, the method seems to work, and, as I mentioned, the authors already include most of the experiments that I'd think to ask for. On the other: Table 5 shows that the lion's share of the improvement from Metalic comes from the in-context meta-learning; the per-task fine-tuning phase only helps in the many-shot case. But Metalic is only the best method in the zero-shot case (albeit with fewer parameters than baselines), where the only possible improvement comes from meta-learning. If the authors provided parameter-matched versions of Metalic that showed a much clearer benefit to using the method even in the many-shot case, that might change my thinking, but that also obviously changes the viability of performing the fine-tuning step at inference time. As it stands, the paper seems to show that in-context meta-learning works well  for zero-shot protein fitness prediction. In-context meta-learning isn't new, and I just don't know if a simple application of a method that already exists to a task that---for ICLR readers---is ultimately fairly niche warrants acceptance. I'm obviously a bit conflicted and can likely be convinced, but at the moment I'm leaning reject.

Meta-review: The authors propose Metalic (Meta-Learning In-Context), a meta-learning framework for protein fitness prediction. By combining in-context learning with fine-tuning, Metalic adapts to new tasks with fewer parameters and better generalization. It outperforms state-of-the-art models in low-data settings, setting a new benchmark on ProteinGym. This work demonstrates the potential of meta-learning in advancing protein engineering, particularly in data-scarce scenarios.

**Strengths:**

- The integration of in-context meta-learning with fine-tuning is novel and computationally efficient compared to state-of-the-art approaches.
- The paper is clearly written and highly reproducible.
- **Metalic** demonstrates strong performance in zero-shot and few-shot protein fitness prediction tasks, which is especially valuable given the scarcity of labeled data in protein-related fields.

**Weaknesses:**

- More experiments and analyses are needed to strengthen the findings.
- The novelty of the approach should be explained in greater detail.

**Overall:** This paper addresses a significant problem in protein learning. While the use of meta-learning in this domain is not entirely new, the specific combination of in-context meta-learning followed by fine-tuning, along with the proposed fine-tuning approach, represents a novel and valuable contribution. Although the results do not achieve state-of-the-art performance across all tasks, they are strong enough to support the authors' claims. I recommend acceptance.

Decision: Accept (Poster)