Aggregating residue-level protein language model embeddings with optimal transport

Area of Science:

• Computational Biology
• Bioinformatics
• Machine Learning

Background:

• Protein language models (PLMs) generate per-residue embeddings from protein sequences.
• Variable-length outputs from PLMs hinder protein-level prediction tasks requiring uniform inputs.
• Average pooling is a common but potentially suboptimal method for summarizing PLM outputs.

Purpose of the Study:

• To develop a novel method for converting variable-length PLM embeddings into fixed-length representations.
• To address the challenge of variable sequence lengths in protein representation learning.
• To improve the performance of protein-level prediction tasks using PLMs.

Main Methods:

• Utilized optimal transport theory to aggregate per-token PLM outputs.
• Conceptualized token embeddings as samples from a probability distribution.
• Employed sliced-Wasserstein distances to create fixed-length, protein-level embeddings.

Main Results:

• The proposed optimal transport method generates fixed-length embeddings independent of protein length.
• This method demonstrated superior performance compared to average pooling across various downstream tasks.
• Smaller PLMs using this method achieved performance comparable to larger PLMs using average pooling, especially for longer sequences.

Conclusions:

• Optimal transport provides an effective strategy for generating fixed-length protein embeddings from PLMs.
• The method enhances the utility of PLMs for protein-level predictions, particularly for long sequences.
• This approach enables more efficient use of smaller PLMs while maintaining high performance.