The CAFA6 Challenge: predicting protein function from sequence

Now that the Arc Virtual Cell Challenge has finished, my attention has shifted to a different type of computational biology-related competition: the CAFA6 challenge, hosted on Kaggle.

The task in the CAFA6 challenge is to predict the function of a protein, starting from its sequence. In practical terms, we need to predict the Gene Ontology terms associated with each protein in the target dataset.

I think this challenge is a good entry point for bioinformaticians wishing to understand more of deep learning. Here I am going to describe my approach, which includes using a foundation model (ESM2) to transform protein sequences into embeddings, using another model (Gemma) to transform Gene Ontology terms into embeddings, and train another model to make a prediction.

This challenge would be very difficult to approach using traditional bioinformatics tools - but with a bit of deep learning, we can give it a try.

The input data for this competition consists of a file containing FASTA sequences (train/test_sequence.fasta), and another containing the GO terms associated with each protein.

There is also a file containing the descriptions of the terms, in the obo format, and another containing the taxonomy - which species does each protein belongs to.

Notably, each protein can be associated with multiple Gene Ontology terms - for example Q5W0B1 in the screenshot above includes multiple rows, one for each term. Moreover, the protein annotations are likely to be incomplete - there will be many proteins that should be associated with gene ontology terms, but they are not in the data because there is no experimental evidence or just nobody has annotated the pair. This is going to make the modelling more difficult.

The task is to predict the terms associated with each protein. For example, for protein Q5W0B1 we know the sequence (in the train_sequences.fasta file) and the associated gene ontology terms. We want to teach a model to learn how to associate protein sequences to terms, and use it to predict the annotations for the proteins in the test dataset.

If you come from a pure computational biology background, like me, you may struggle to approach this challenge. The only input we have is the protein sequence - I can’t really think of an approach to do this, without using a bit of deep learning. Maybe we could use an external database to get information, given the protein ids? But that would defeat the purpose of the challenge.

My plan for this competition is relatively simple. It is based on an approach called “Two-Towers model”, traditionally used in recommendation systems.

In short, we use two networks: one for protein sequences and one for Gene Ontology terms. Each network converts its input into a numerical representation, called an embedding. In this way, both sequences and terms are mapped into a shared numerical space. We then train an additional model to learn how to associate each protein embedding with its corresponding GO term embeddings.

The first step is to convert the protein sequences into embeddings.

These embeddings should not be random numbers, but they should capture the structure, conservation, and function of the protein sequences. If two proteins are similar in terms of sequence, their distance in the embedding space should be small. Proteins derived from the same gene, or from paralogs, should also be similar to each other. How can we compute such numbers?

Likely, there are a few foundation models available in the literature, that can do just that. In my case, I am using the ESM2 (or the ESM-Cambrian) model, trained by Evolutionary Scale. This model has been trained on a large collection of sequences, and it is freely available on Hugging Face

This is my notebook for the conversion. Essentially, I just download the checkpoint from Hugging Face, and use it to compute the embeddings for all the proteins in the train and test dataset. You are welcome to use the dataset, if you like to save a few hours of compute.

The second step is to convert the Gene Ontology terms into embeddings.

I’ve seen many strategies to deal with Gene Ontology terms. They are always difficult to handle - because of their hierarchical nature, their incompleteness, their difficult ontology.

So, this time, I want to try an approach used by the CyteOnto package from Nygen, described in a previous article here.

Essentially, we ask an LLM, like Gemma, to describe each Gene Ontology term; then we take the embeddings for these descriptions, also from Gemma.

This approach transforms the original Gene Ontology terms, like GO:0070902, into vectors of numbers, in a similar way as the protein embeddings. Terms that have a similar meaning for Gemma should be close to each other in this space; for example, “Intrinsic to the luminal side of the endoplasmic reticulum” should be close to “Intrinsic to the cytoplasmic side of the endoplasmic reticulum”, but also closer to other luminal terms compared to the other term.

Here is my notebook on Kaggle. I am still refining it. It uses the Gemma model from Kaggle, loading it using the transformers library - there is no need for any API key.

The next step is to align these two sets of embeddings together.

I haven’t got to this point yet - it is my plan for the Xmas holidays, if I can afford it.

In a standard dual-tower setup, the two encoders are trained jointly, with similarity between embeddings (often cosine similarity) used as the training signal. Here, however, I start from pre-computed embeddings, which leads to a slightly different architecture.

There are many challenges in this alignment step. Each protein can have multiple correct GO terms (multi-label), many terms are rare (class imbalance), and the labels are incomplete - an absent annotation is not a true negative. On top of that, GO is hierarchical, so we want the model to learn that predicting a specific child term should also imply its parents.

My plan is to keep the foundation embeddings frozen at first and learn a lightweight “projection head” that maps proteins and GO terms into a shared space. But, this is not going to be straightforward to implement.

Fancy to join me on this challenge on Kaggle? There are about two months left before the deadline. Please follow this substack, and I’ll keep you posted on the progress.