Scaling down for efficiency: Medium-sized protein language models perform well at transfer learning on realistic datasets

  1. Luiz C. Vieira
  2. Morgan L. Handojo
  3. Claus O. Wilke

Reviewed by

Read the full article See related articles

Listed in

Log in to save this article

Abstract

Protein language models (pLMs) can offer deep insights into evolutionary and structural properties of proteins. While larger models, such as the 15 billion parameter model ESM-2, promise to capture more complex patterns in sequence space, they also present practical challenges due to their high dimensionality and high computational cost. We systematically evaluated the performance of various pLMs across multiple biological datasets to assess the impact of model size on transfer learning. Surprisingly, we found that larger models not necessarily outperform smaller ones, in particular when data is limited. Medium-sized models, such as ESM-2 650M and ESM C 600M, demonstrated consistently good performance, falling only slightly behind their larger counterparts—ESM-2 15B and ESM C 6B—despite being many times smaller. Additionally, we compared various methods of compressing embeddings prior to transfer learning, and we found that mean embeddings consistently outperformed other compression methods. In summary, ESM C 600M with mean embeddings offers an optimal balance between performance and efficiency, making it a practical and scalable choice for transfer learning in realistic biological applications.

This work challenges the common belief that larger language models always yield better results, here in the context of protein biochemistry. By systematically comparing transformer models of different sizes in transfer learning tasks, we demonstrate that medium size models, such as ESM C 600M, frequently perform as well as or better than larger variants, especially when data is limited. These findings provide an efficient strategy for machine learning-based protein analysis. Smaller and more efficient models help democratize cutting-edge AI approaches, making them more accessible to researchers with limited computational resources.

Article activity feed

  1. Version published to 10.1101/2024.11.22.624936v2 on bioRxiv
  2. Arcadia Science

    We have evaluated the performance of ESM2 embeddings across various model sizes (from 8 million to 15 billion parameters) in transfer learning tasks on a wide range of different biological datasets

    I think the diversity of regression tasks lends a lot of robustness to your conclusions. However, I think you're using the term "transfer learning" rather narrowly, specifically referring to prediction tasks where either a value or a vector is predicted for each sequence.

    There are many classes of transfer learning tasks, like sequence labeling, token classification, all sequence-to-sequence tasks, etc. I think being more specific about the type of transfer learning you guys are making claims about would make your conclusions more accurate.

  3. Arcadia Science

    Scaling Down for Efficiency: Medium-Sized Transformer Models for Protein Sequence Transfer Learning

    Thanks for this insightful piece. I've left some food for thought below.

  4. Arcadia Science

    Even though these models were also pretrained with a maximum sequence length

    Technically ESM2 is trained using sequences longer than 1022, but a length 1022 subsequence is sampled whenever it is selected for a training batch.

  5. Arcadia Science

    Mean reduction in R2 when embeddings are compressed with methods other than mean pooling.A) Results for DMS data. B) Results for diverse protein sequences (PISCES data). In all cases, the y-axis represents different compression methods and the x-axis shows the resulting difference in R2. Dots represent the fixed effects estimates from mixed-effects modeling, and error bars represent 95% confidence intervals.

    This analysis that compares pooling methods was very informative, but it left me wondering the extent that mean pooling compares to no pooling at all. Is this something y'all considered? It would be interesting to compare the R2 of a more sophisticated transfer learning model that ingests the raw embeddings (like a basic FCN). Though an apples to apples might be hard to create, it would be useful to know the "cost" of mean pooling by observing the extend to which raw embeddings outperform mean pooling (if at all?)

  6. Arcadia Science

    In most scenarios

    I really don't think this is true. Many transfer learning tasks are token-level predictions, and therefore in those scenarios embeddings cannot be compressed.

  9. Version published to 10.1101/2024.11.22.624936v1 on bioRxiv