![\"\"](\"https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F9675%2Fc5eb213eac49de03a7fcc9f51dcae0ca%2Fmodel_architecture.png?generation=1790505060872179&alt=media\")
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.
\nThe source code for the solution and instructions how to reproduce it are available in this repository: https://github.com/stefanistefanov/CAFA6-Solution. The input datasets are provided as Kaggle datasets and referenced below.
\nThe inputs to the models are as follows:
\nfacebook/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-3diNeuML/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-pubmedembeddings directory. The script with commands for generating them is create_embeddings.sh in the repository.test_terms_evidence.tsv file).ete3.NCBITaxa from the NCBI taxonomy database. Biological Taxonomy post by @mtinti
The model architecture fuses the heterogeneous inputs above using several components:
\nTermEvidenceEncoder — a 2-layer Conv1D network projecting the term-evidence one-hot encoded matrix to term logits. The idea is from @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
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.
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.
To make the model more robust, the following data augmentations are applied during training and prediction:
\nMultiple 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.
Heinzinger, M. et al. (2023). ProstT5: Bilingual Language Model for Protein Sequence and Structure. bioRxiv. doi: 10.1101/2023.07.23.550085.
\nJumper, J. et al. (2021). Highly accurate protein structure prediction with AlphaFold. Nature, 596, pages 583–589. doi: 10.1038/s41586-021-03819-2.
\nKempen, M. et al. (2022). Fast and accurate protein structure search with Foldseek. bioRxiv. doi: 10.1101/2022.02.07.479398
\nLin, Z. et al. (2022). Evolutionary-scale prediction of atomic level protein structure with a language model. bioRxiv. doi: 10.1101/2022.07.20.500902.
\nMezzetti, D. (2023). PubMedBERT Embeddings Model. NeuML. https://huggingface.co/NeuML/pubmedbert-base-embeddings
\nmtinti. Biological Taxonomy post. https://www.kaggle.com/competitions/cafa-6-protein-function-prediction/discussion/613750. 2025. Kaggle
\ntito. 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. 2023. Kaggle
\nVaradi, 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
", "messages": [ { "id": "3480041", "postDate": "06/24/2026 11:51:13", "content": "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.
\nThe source code for the solution and instructions how to reproduce it are available in this repository: https://github.com/stefanistefanov/CAFA6-Solution. The input datasets are provided as Kaggle datasets and referenced below.
\nThe inputs to the models are as follows:
\nfacebook/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-3diNeuML/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-pubmedembeddings directory. The script with commands for generating them is create_embeddings.sh in the repository.test_terms_evidence.tsv file).ete3.NCBITaxa from the NCBI taxonomy database. Biological Taxonomy post by @mtinti
The model architecture fuses the heterogeneous inputs above using several components:
\nTermEvidenceEncoder — a 2-layer Conv1D network projecting the term-evidence one-hot encoded matrix to term logits. The idea is from @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
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.
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.
To make the model more robust, the following data augmentations are applied during training and prediction:
\nMultiple 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.
Heinzinger, M. et al. (2023). ProstT5: Bilingual Language Model for Protein Sequence and Structure. bioRxiv. doi: 10.1101/2023.07.23.550085.
\nJumper, J. et al. (2021). Highly accurate protein structure prediction with AlphaFold. Nature, 596, pages 583–589. doi: 10.1038/s41586-021-03819-2.
\nKempen, M. et al. (2022). Fast and accurate protein structure search with Foldseek. bioRxiv. doi: 10.1101/2022.02.07.479398
\nLin, Z. et al. (2022). Evolutionary-scale prediction of atomic level protein structure with a language model. bioRxiv. doi: 10.1101/2022.07.20.500902.
\nMezzetti, D. (2023). PubMedBERT Embeddings Model. NeuML. https://huggingface.co/NeuML/pubmedbert-base-embeddings
\nmtinti. Biological Taxonomy post. https://www.kaggle.com/competitions/cafa-6-protein-function-prediction/discussion/613750. 2025. Kaggle
\ntito. 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. 2023. Kaggle
\nVaradi, 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
", "rawMarkdown": "## Overview\nThe 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.\n\nThe source code for the solution and instructions how to reproduce it are available in this repository: https://github.com/stefanistefanov/CAFA6-Solution. The input datasets are provided as Kaggle datasets and referenced below.\n\n\n## Model Inputs\nThe inputs to the models are as follows:\n- `facebook/esm2_t36_3B_UR50D` model embeddings of the protein sequences.\n- `Rostlab/ProstT5_fp16` model embeddings of the protein sequences.\n- `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\n- `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\n- The generated embeddings are available at 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.\n- 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).\n- 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\n- The normalized sequence length — the sequence length divided by 10000 and clipped to 1.\n\n\n## Model Architecture\n\n\nThe model architecture fuses the heterogeneous inputs above using several components:\n- **TermEvidenceEncoder** — a 2-layer Conv1D network projecting the term-evidence one-hot encoded matrix to term logits. The idea is from @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\n\n- **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.\n\n- **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.\n\n- 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.\n\n- 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.\n\n- 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.\n\n## Training\nThe script with commands for training the models is `train_models.sh` in the repository.\n- 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.\n\n- 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.\n\n- 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.\n\n- 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.\n\n- 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.\n\n- 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.\n\n- 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.\n\n- To optimize the CAFA metric computations for speed and memory, they are implemented with sparse matrices.\n\n## Data Augmentations\nTo make the model more robust, the following data augmentations are applied during training and prediction:\n- Input GO terms with non-experimental evidence are randomly dropped (replaced with an unknown input term).\n- To simulate unseen taxons, more specific taxons in the lineage are randomly replaced with an unknown taxon indicator.\n- Whole input feature types (a given embedding type, the input GO terms, or the taxon's lineage) are randomly dropped.\n\n## Ensembling \nMultiple 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`.\n- 10 predictions are generated from models 1, 2, and 3. For each model, one of the predictions is generated without data augmentations.\n- 5 predictions are generated from model 4. One of these predictions is generated without data augmentations.\n- The trained models for the ensemble are available at https://www.kaggle.com/datasets/stefanstefanov/cafa6-models \n\n## References\n\nHeinzinger, M. et al. (2023). ProstT5: Bilingual Language Model for Protein Sequence and Structure. bioRxiv. doi: 10.1101/2023.07.23.550085.\n\nJumper, J. et al. (2021). Highly accurate protein structure prediction with AlphaFold. Nature, 596, pages 583–589. doi: 10.1038/s41586-021-03819-2.\n\nKempen, M. et al. (2022). Fast and accurate protein structure search with Foldseek. bioRxiv. doi: 10.1101/2022.02.07.479398\n\nLin, Z. et al. (2022). Evolutionary-scale prediction of atomic level protein structure with a language model. bioRxiv. doi: 10.1101/2022.07.20.500902.\n\nMezzetti, D. (2023). PubMedBERT Embeddings Model. NeuML. https://huggingface.co/NeuML/pubmedbert-base-embeddings\n\nmtinti. Biological Taxonomy post. https://www.kaggle.com/competitions/cafa-6-protein-function-prediction/discussion/613750. 2025. Kaggle\n\ntito. 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. 2023. Kaggle\n\nVaradi, 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", "votes": null }, { "id": "3486401", "postDate": "07/02/2026 14:00:39", "content": "Great work ! \nOn CAFA5 did you use the same approach ?
", "rawMarkdown": "Great work ! \nOn CAFA5 did you use the same approach ?", "votes": null }, { "id": "3487969", "postDate": "07/04/2026 12:39:08", "content": "Thanks! Overall, the CAFA 6 approach is similar to the one from CAFA 5. I liked tito's \n3rd place solution idea for using non-experimental annotations and included a TermEvidenceEncoder in CAFA 6 based on it. Additionally, I introduced a terms logits gate to control the sum, and a TermEvidenceEmbeddingEncoder to have non-experimental term evidence embeddings as an early input. For CAFA 6, the taxon's lineage is encoded with a transformer encoder layer, while for CAFA 5 only a taxon id embedding was used. Structure-related representation with Foldseek 3Di is also newly added for CAFA 6, while for CAFA 5 no structure information was included.
Great work ! \nOn CAFA5 did you use the same approach ?
", "votes": null, "replies": [ { "id": 3487969, "author_name": "stefanstefanov", "author_url": "", "post_date": "07/04/2026 12:39:08", "content": "Thanks! Overall, the CAFA 6 approach is similar to the one from CAFA 5. I liked tito's \n3rd place solution idea for using non-experimental annotations and included a TermEvidenceEncoder in CAFA 6 based on it. Additionally, I introduced a terms logits gate to control the sum, and a TermEvidenceEmbeddingEncoder to have non-experimental term evidence embeddings as an early input. For CAFA 6, the taxon's lineage is encoded with a transformer encoder layer, while for CAFA 5 only a taxon id embedding was used. Structure-related representation with Foldseek 3Di is also newly added for CAFA 6, while for CAFA 5 no structure information was included.