# 6th Place Solution for the CAFA 6 Protein Function Prediction

Overview
The solution is an ensemble of 4 neural network models with the same architecture, combining heterogeneous inputs, and producing probability prediction for each GO term. The inputs vary across the models and include sequence, structure, and text embeddings, taxon's lineage, non-experimental annotations, and evidence codes.
The source code for the solution and instructions how to reproduce it are available in this repository:
[https://github.com/stefanistefanov/CAFA6-Solution](https://github.com/stefanistefanov/CAFA6-Solution)
. The input datasets are provided as Kaggle datasets and referenced below.
Model Inputs
The inputs to the models are as follows:
facebook/esm2_t36_3B_UR50D
model embeddings of the protein sequences.
Rostlab/ProstT5_fp16
model embeddings of the protein sequences.
Rostlab/ProstT5_fp16
model embeddings of the 3Di sequences generated from AlphaFoldDB 3D structures using foldseek. Dataset:
[https://www.kaggle.com/datasets/stefanstefanov/cafa6-3di](https://www.kaggle.com/datasets/stefanstefanov/cafa6-3di)
NeuML/pubmedbert-base-embeddings
model embeddings of text concatenated from the article title, abstract, chemicals, and MeSH headings downloaded from PubMed using Entrez with the PubMed IDs from the UniProt ID mapping file. The 5 most recent articles are used, and the mean and max of their embeddings are concatenated and used as model input. Dataset:
[https://www.kaggle.com/datasets/stefanstefanov/cafa6-pubmed](https://www.kaggle.com/datasets/stefanstefanov/cafa6-pubmed)
The generated embeddings are available at
[https://www.kaggle.com/datasets/stefanstefanov/cafa6-data](https://www.kaggle.com/datasets/stefanstefanov/cafa6-data)
under the
embeddings
directory. The script with commands for generating them is
create_embeddings.sh
in the repository.
GO terms with non-experimental evidence('IBA', 'IBD', 'IEA', 'IGC', 'IKR', 'IRD', 'ISA', 'ISM', 'ISO', 'ISS', 'NAS', 'ND', 'RCA') and their evidence codes are used as inputs (
test_terms_evidence.tsv
file).
The taxon's lineage, retrieved with
ete3.NCBITaxa
from the NCBI taxonomy database.
[Biological Taxonomy](https://www.kaggle.com/competitions/cafa-6-protein-function-prediction/discussion/613750)
post by
[@mtinti](/mtinti)
The normalized sequence length — the sequence length divided by 10000 and clipped to 1.
Model Architecture
![](assets/d6ca127c139ef3ffdc5d6a39.png)
The model architecture fuses the heterogeneous inputs above using several components:
TermEvidenceEncoder
— a 2-layer Conv1D network projecting the term-evidence one-hot encoded matrix to term logits. The idea is from
[@its7171](/its7171)
3rd Place Solution for CAFA 5:
[https://www.kaggle.com/competitions/cafa-5-protein-function-prediction/writeups/tito-3rd-place-solution-for-the-cafa-5-protein-fun](https://www.kaggle.com/competitions/cafa-5-protein-function-prediction/writeups/tito-3rd-place-solution-for-the-cafa-5-protein-fun)
TermEvidenceEmbeddingEncoder
— a single-layer transformer encoder over the sum of the term and evidence embeddings, producing a mean-pooled representation (dim 512) of the term-evidence associations. Per protein, the non-experimental GO terms and their respective evidence codes are passed as two corresponding sequences to this encoder. The term and evidence embeddings are learned during training.
TaxonEmbeddingEncoder
— a single-layer transformer encoder over the taxon lineage sequence, producing a mean-pooled representation (dim 512) of the protein's taxon lineage. The taxon embeddings are learned during training.
A single feedforward layer per input embedding (sequence/structure/PubMed embedding) reduces them to the same dimension of 1024. The reduced embeddings, the normalized sequence length, and the TermEvidenceEmbeddingEncoder and TaxonEmbeddingEncoder output representations are concatenated into a single feature vector.
There is a separate head for each of the three sub-ontologies (BP, CC, MF), consisting of one feedforward layer with a GELU non-linearity followed by a Dropout with 0.5 probability and an output linear layer producing per-term logits for the respective sub-ontology. The hidden dimension of the feedforward layer differs for each sub-ontology (3072/1024/1536 for BP/CC/MF respectively) as they have different numbers of terms. The feature vector is fed to each of the three heads.
A logits gate is introduced in each head; it controls how much the term logits from the TermEvidenceEncoder influence the final term prediction. The feature vector is also fed to the logits gate to provide a signal for the gate's behavior.
Training
The script with commands for training the models is
train_models.sh
in the repository.
Model 1 is trained on esm2_t36_3B_UR50D sequence embeddings, ProstT5 3Di embeddings, PubMed embeddings, and GO term annotations with non-experimental evidence codes from goa_uniprot_all.gaf release 229 from December 2025.
Model 2 is trained on esm2_t36_3B_UR50D sequence embeddings and GO term annotations with non-experimental evidence codes from goa_uniprot_all.gaf release 226 from May 2025. It is trained on a larger set of proteins from both Swiss-Prot and TrEMBL from January 2026.
Model 3 is trained on esm2_t36_3B_UR50D and ProstT5 sequence embeddings, ProstT5 3Di embeddings, and GO term annotations with non-experimental evidence codes from goa_uniprot_all.gaf release 229 from December 2025.
Model 4 is trained on esm2_t36_3B_UR50D and ProstT5 sequence embeddings, ProstT5 3Di embeddings, PubMed embeddings, and GO term annotations with non-experimental evidence codes from goa_uniprot_all.gaf release 229 from December 2025.
All 4 models are trained for 15 epochs, and the checkpoint of the last epoch is used for submission. The AdamW optimizer with weight decay 0.01 and a OneCycleLR scheduler are used. The training batch size is 128. The loss function is BCEWithLogitsLoss with label weights computed as the square root of the IA weights and capped at 0.5.
While the submission models are trained on release 229, the latest available at that time, model development and parameter tuning(such as the number of training epochs) were performed on the earlier release 226, using release 229 for the validation set.
Training on the earlier release 226 and validating on the latest release 229 enabled temporal validation, simulating the prediction of future annotations and allowing to compute the CAFA6 metrics on no-knowledge, limited-knowledge, and partial-knowledge protein targets.
To optimize the CAFA metric computations for speed and memory, they are implemented with sparse matrices.
Data Augmentations
To make the model more robust, the following data augmentations are applied during training and prediction:
Input GO terms with non-experimental evidence are randomly dropped (replaced with an unknown input term).
To simulate unseen taxons, more specific taxons in the lineage are randomly replaced with an unknown taxon indicator.
Whole input feature types (a given embedding type, the input GO terms, or the taxon's lineage) are randomly dropped.
Ensembling
Multiple predictions are generated with each of the models using different random seeds and data augmentations enabled during predict. The final submission is an average of 35 such predictions. The script with commands for generating the predictions and ensembling is
predict_ensemble.sh
.
10 predictions are generated from models 1, 2, and 3. For each model, one of the predictions is generated without data augmentations.
5 predictions are generated from model 4. One of these predictions is generated without data augmentations.
The trained models for the ensemble are available at
[https://www.kaggle.com/datasets/stefanstefanov/cafa6-models](https://www.kaggle.com/datasets/stefanstefanov/cafa6-models)
References
Heinzinger, M. et al. (2023). ProstT5: Bilingual Language Model for Protein Sequence and Structure. bioRxiv. doi: 10.1101/2023.07.23.550085.
Jumper, J. et al. (2021). Highly accurate protein structure prediction with AlphaFold. Nature, 596, pages 583–589. doi: 10.1038/s41586-021-03819-2.
Kempen, M. et al. (2022). Fast and accurate protein structure search with Foldseek. bioRxiv. doi: 10.1101/2022.02.07.479398
Lin, Z. et al. (2022). Evolutionary-scale prediction of atomic level protein structure with a language model. bioRxiv. doi: 10.1101/2022.07.20.500902.
Mezzetti, D. (2023). PubMedBERT Embeddings Model. NeuML.
[https://huggingface.co/NeuML/pubmedbert-base-embeddings](https://huggingface.co/NeuML/pubmedbert-base-embeddings)
mtinti. Biological Taxonomy post.
[https://www.kaggle.com/competitions/cafa-6-protein-function-prediction/discussion/613750](https://www.kaggle.com/competitions/cafa-6-protein-function-prediction/discussion/613750)
. 2025. Kaggle
tito. 3rd Place Solution for the CAFA 5 Protein Function Prediction.
[https://www.kaggle.com/competitions/cafa-5-protein-function-prediction/writeups/tito-3rd-place-solution-for-the-cafa-5-protein-fun](https://www.kaggle.com/competitions/cafa-5-protein-function-prediction/writeups/tito-3rd-place-solution-for-the-cafa-5-protein-fun)
. 2023. Kaggle
Varadi, M. et al. (2023). AlphaFold Protein Structure Database in 2024: Providing structure coverage for over 214 million protein sequences. Nucleic Acids Research, gkad1011. doi: 10.1093/nar/gkad1011