This competition was also the focus of my BTech Major Project and I have invested approximately 75 hours into this competition. The goal of the CAFA 5 competition is to predict biological function of proteins given their amino-acid sequences. One interesting part about the competition is that models are evaluated on \"future data,\" meaning the test set does not exist at submission and the ground truth is compiled months later from newly published lab results.
\nI had two main strategies for winning this competition:
\nKaggle Notebook: CAFA ESM JAX 15B
\nIn Kaggle you can get either 30 hours of the P100 GPU or 20 hours of 8xTPUv3 for running modelling experiments. Simple math says that 20 hours of 8 TPUv3 will give you 13.2x more FLOPs than 30 hours of P100 GPUs. So there is a huge competitive advantage in being able to use TPUs effectively as compared to using consumer grade GPUs.
\nRight now ESM2 (15B params) is the largest and the best protein language model available. However, it does not fit in a single TPU or a GPU so I rewrote ESM2 in JAX with GSPMD-style tensor partitioning to fit the model across the 8 TPUs.
\nEmbedding Strategy: The main goal is to get the best possible embedding for a protein. I tried the last layer CLS embedding and also tried mean pooling to get 5k dim vector for every protein. Training an MLP on these embeddings gave me a score of ~0.54.
\nOne failed experiment was that I tried to train a secondary transformer on mean embeddings per amino acid, but this did not outperform the baseline MLP.
\nKaggle Notebook (Very Rough State)
\nTail Class Problem: There are approximately 40k unique GO terms. Many specific functions (leaf nodes) have very few training samples. My hypothesis was that we could enrich these \"tail\" terms by aggregating information from their more frequent \"parent\" and \"sibling\" nodes.
\nMy second secret strategy was using GNNs to build better protein embeddings by using the geometric structural dependencies within the Gene Ontology (GO) hierarchy.
\nNow instead of treating every protein as independent and just using the amino acid sequence of the protein to build their embedding, we can now also use these parent, child, and sibling relationships between the proteins to build even better protein embeddings.
\nI did not see anyone in this competition using this Gene Ontology structure so I took this as another competitive advantage. Using GNNs was also one of my main reasons for entering this competition since I have gotten very interested in Graph Neural Networks recently.
\nTo do this I focused on the top 3k term vectors and built a 3000x3000 adjacency matrix. Then I used power iteration to build a transition matrix:
\nMy code here was brute force but thanks to the power of TPUs it ended up working anyways. I also tried spectral position embedding but it did not give a better score.
\nFull Finetuning of the ProtT5 model: I fully finetuned the ProtT5 model on TPUs. This also did not give a better score as compared to just using fixed embedding and training an MLP on it.
\nDifferent Loss Functions: To address label imbalance, I tried different loss functions including focal loss, dice loss, poly loss, and the classic approach of up-weighting low-sample terms. All of these losses gave very similar scores which was a little surprising as I thought this would play a bigger role.
", "messages": [ { "id": 2402415, "postDate": "2023-08-22T06:23:40.867Z", "content": "This competition was also the focus of my BTech Major Project and I have invested approximately 75 hours into this competition. The goal of the CAFA 5 competition is to predict biological function of proteins given their amino-acid sequences. One interesting part about the competition is that models are evaluated on \"future data,\" meaning the test set does not exist at submission and the ground truth is compiled months later from newly published lab results.
\nI had two main strategies for winning this competition:
\nKaggle Notebook: CAFA ESM JAX 15B
\nIn Kaggle you can get either 30 hours of the P100 GPU or 20 hours of 8xTPUv3 for running modelling experiments. Simple math says that 20 hours of 8 TPUv3 will give you 13.2x more FLOPs than 30 hours of P100 GPUs. So there is a huge competitive advantage in being able to use TPUs effectively as compared to using consumer grade GPUs.
\nRight now ESM2 (15B params) is the largest and the best protein language model available. However, it does not fit in a single TPU or a GPU so I rewrote ESM2 in JAX with GSPMD-style tensor partitioning to fit the model across the 8 TPUs.
\nEmbedding Strategy: The main goal is to get the best possible embedding for a protein. I tried the last layer CLS embedding and also tried mean pooling to get 5k dim vector for every protein. Training an MLP on these embeddings gave me a score of ~0.54.
\nOne failed experiment was that I tried to train a secondary transformer on mean embeddings per amino acid, but this did not outperform the baseline MLP.
\nKaggle Notebook (Very Rough State)
\nTail Class Problem: There are approximately 40k unique GO terms. Many specific functions (leaf nodes) have very few training samples. My hypothesis was that we could enrich these \"tail\" terms by aggregating information from their more frequent \"parent\" and \"sibling\" nodes.
\nMy second secret strategy was using GNNs to build better protein embeddings by using the geometric structural dependencies within the Gene Ontology (GO) hierarchy.
\nNow instead of treating every protein as independent and just using the amino acid sequence of the protein to build their embedding, we can now also use these parent, child, and sibling relationships between the proteins to build even better protein embeddings.
\nI did not see anyone in this competition using this Gene Ontology structure so I took this as another competitive advantage. Using GNNs was also one of my main reasons for entering this competition since I have gotten very interested in Graph Neural Networks recently.
\nTo do this I focused on the top 3k term vectors and built a 3000x3000 adjacency matrix. Then I used power iteration to build a transition matrix:
\nMy code here was brute force but thanks to the power of TPUs it ended up working anyways. I also tried spectral position embedding but it did not give a better score.
\nFull Finetuning of the ProtT5 model: I fully finetuned the ProtT5 model on TPUs. This also did not give a better score as compared to just using fixed embedding and training an MLP on it.
\nDifferent Loss Functions: To address label imbalance, I tried different loss functions including focal loss, dice loss, poly loss, and the classic approach of up-weighting low-sample terms. All of these losses gave very similar scores which was a little surprising as I thought this would play a bigger role.
", "rawMarkdown": "This competition was also the focus of my BTech Major Project and I have invested approximately 75 hours into this competition. The goal of the CAFA 5 competition is to predict biological function of proteins given their amino-acid sequences. One interesting part about the competition is that models are evaluated on \"future data,\" meaning the test set does not exist at submission and the ground truth is compiled months later from newly published lab results.\n\nI had two main strategies for winning this competition: \n1. **TPU-Maxxing**: Rewriting the ESM2-15B model in JAX to utilize TPUv3-8 pods via GSPMD tensor partitioning.\n2. **GNNs**: Using the Gene Ontology (GO) Hierarchy through Geometric Deep Learning techniques to improve embeddings for sparse \"tail\" classes.\n\n---\n\n## Scaling up to ESM-15B with TPUs\n\n[Kaggle Notebook: CAFA ESM JAX 15B](https://www.kaggle.com/code/readoc/cafa-esm-jax-15b/edit/run/133312976)\n\nIn Kaggle you can get either 30 hours of the P100 GPU or 20 hours of 8xTPUv3 for running modelling experiments. Simple math says that 20 hours of 8 TPUv3 will give you 13.2x more FLOPs than 30 hours of P100 GPUs. So there is a huge competitive advantage in being able to use TPUs effectively as compared to using consumer grade GPUs.\n\nRight now [ESM2](https://www.biorxiv.org/content/10.1101/2022.07.20.500902v2.full.pdf) (15B params) is the largest and the best protein language model available. However, it does not fit in a single TPU or a GPU so I rewrote ESM2 in JAX with [GSPMD](https://arxiv.org/pdf/2105.04663)-style tensor partitioning to fit the model across the 8 TPUs.\n\n**Embedding Strategy**: The main goal is to get the best possible embedding for a protein. I tried the last layer CLS embedding and also tried mean pooling to get 5k dim vector for every protein. Training an MLP on these embeddings gave me a score of ~0.54.\n\nOne failed experiment was that I tried to train a secondary transformer on mean embeddings per amino acid, but this did not outperform the baseline MLP.\n\n---\n\n## GO Hierarchy Aggregation\n\n[Kaggle Notebook (Very Rough State)](https://www.kaggle.com/code/readoc/cafa-keras-mlp-gnn-v2)\n\n**Tail Class Problem**: There are approximately 40k unique GO terms. Many specific functions (leaf nodes) have very few training samples. My hypothesis was that we could enrich these \"tail\" terms by aggregating information from their more frequent \"parent\" and \"sibling\" nodes.\n\nMy second secret strategy was using GNNs to build better protein embeddings by using the geometric structural dependencies within the Gene Ontology (GO) hierarchy. \n\n### Example Relationship between Proteins\n\n* **Protein A (Child Node)**: Annotated with a very specific rare leaf node. It is a tail class with low sample frequency.\n* **Protein B (Parent Node)**: Represents a high-level biological process. It has a high sample frequency in the dataset.\n* **Protein C (Sibling Node to Protein A)**: Another specialized sub-process sharing the same parent (Protein B) as Protein A.\n\nNow instead of treating every protein as independent and just using the amino acid sequence of the protein to build their embedding, we can now also use these parent, child, and sibling relationships between the proteins to build even better protein embeddings. \n\nI did not see anyone in this competition using this Gene Ontology structure so I took this as another competitive advantage. Using GNNs was also one of my main reasons for entering this competition since I have gotten very interested in Graph Neural Networks recently.\n\n### Neighborhood Aggregation Logic\n\nTo do this I focused on the top 3k term vectors and built a 3000x3000 adjacency matrix. Then I used power iteration to build a transition matrix:\n* **Child Influence**: 0.75 weight (Child -> Parent)\n* **Parent Influence**: 0.25 weight (Parent -> Child)\n\nMy code here was brute force but thanks to the power of TPUs it ended up working anyways. I also tried spectral position embedding but it did not give a better score. \n\n---\n\n## Other Tricks\n\n**Full Finetuning of the ProtT5 model**: I fully finetuned the ProtT5 model on TPUs. This also did not give a better score as compared to just using fixed embedding and training an MLP on it.\n\n**Different Loss Functions**: To address label imbalance, I tried different loss functions including focal loss, dice loss, poly loss, and the classic approach of up-weighting low-sample terms. All of these losses gave very similar scores which was a little surprising as I thought this would play a bigger role.", "votes": 9 }, { "id": 2402552, "postDate": "2023-08-22T07:58:14.647Z", "content": "How can you tell than ESM 15B didn't work based on a score of 0.54? That seems like a strong public LB for a single model.
", "rawMarkdown": "How can you tell than ESM 15B didn't work based on a score of 0.54? That seems like a strong public LB for a single model.", "votes": 4, "replies": [ { "id": 2403865, "postDate": "2023-08-22T22:42:47.467Z", "content": "My thoughts exactly. Not sure that I would have spent so much time to get ESM2 15B to work, as I was getting ~0.52 with ProtT5 alone.
", "rawMarkdown": "My thoughts exactly. Not sure that I would have spent so much time to get ESM2 15B to work, as I was getting ~0.52 with ProtT5 alone.", "votes": 2 } ] }, { "id": 2403867, "postDate": "2023-08-22T22:44:29.320Z", "content": "Surprised that your approach #3 didn't work better.
\nIn my hands focal loss worked better than binary cross-entropy.
", "rawMarkdown": "Surprised that your approach #3 didn't work better.\n\nIn my hands focal loss worked better than binary cross-entropy.", "votes": 1, "replies": [ { "id": 2404063, "postDate": "2023-08-23T03:24:31.657Z", "content": "There was some improvement, but not enough to win the competition
", "rawMarkdown": "There was some improvement, but not enough to win the competition" } ] }, { "id": 2402948, "postDate": "2023-08-22T12:16:36.810Z", "content": "@readoc Thank you so much for sharing, most of the time we can learn more from the things that didn't work than from the things that we know will work.
\n.1. Do you have a link to the repo housing this code, or the notebook that you did this in? Would be curious to take a look
.4. What were your findings based on experimenting with the different loss functions?
", "rawMarkdown": "@readoc Thank you so much for sharing, most of the time we can learn more from the things that didn't work than from the things that we know will work.\n\n.1. Do you have a link to the repo housing this code, or the notebook that you did this in? Would be curious to take a look\"Making this notebook public will be locked until the competition ends\"
\nI am getting this message when trying to make the notebook public
Yeah, unfortunately I got the same error with my notebook :/
", "rawMarkdown": "Yeah, unfortunately I got the same error with my notebook :/", "votes": 1 } ] } ] }, { "id": 2402519, "postDate": "2023-08-22T07:33:45.707Z", "content": "Wow! You rewrote ESM2 in JAX?! If that can save appreciable time or GPU memory, I think it will be very helpful! Did you have a benchmark result vs. pytorch-ESM2?
", "rawMarkdown": "Wow! You rewrote ESM2 in JAX?! If that can save appreciable time or GPU memory, I think it will be very helpful! Did you have a benchmark result vs. pytorch-ESM2?", "votes": 1, "replies": [ { "id": 2404068, "postDate": "2023-08-23T03:27:29.300Z", "rawMarkdown": "", "votes": 1, "isDeleted": true } ] }, { "id": 2402477, "postDate": "2023-08-22T07:02:57.993Z", "content": "Thanks for sharing, but I can't agree with the fourth point. I have similarly tried the fourth point you mentioned and this way can enhance the CV and LB scores and the enhancement is greater in BPO.
", "rawMarkdown": "Thanks for sharing, but I can't agree with the fourth point. I have similarly tried the fourth point you mentioned and this way can enhance the CV and LB scores and the enhancement is greater in BPO.", "votes": 1 } ], "comments": [ { "id": 2402552, "author_name": "John Mitchell", "author_url": "", "post_date": "2023-08-22T07:58:14.647000", "content": "How can you tell than ESM 15B didn't work based on a score of 0.54? That seems like a strong public LB for a single model.
", "votes": 4, "replies": [ { "id": 2403865, "author_name": "Tilii", "author_url": "", "post_date": "2023-08-22T22:42:47.467000", "content": "My thoughts exactly. Not sure that I would have spent so much time to get ESM2 15B to work, as I was getting ~0.52 with ProtT5 alone.
", "votes": 2, "replies": [] } ] }, { "id": 2403867, "author_name": "Tilii", "author_url": "", "post_date": "2023-08-22T22:44:29.320000", "content": "Surprised that your approach #3 didn't work better.
\nIn my hands focal loss worked better than binary cross-entropy.
", "votes": 1, "replies": [ { "id": 2404063, "author_name": "Sarthak Bhatt", "author_url": "", "post_date": "2023-08-23T03:24:31.657000", "content": "There was some improvement, but not enough to win the competition
", "votes": 0, "replies": [] } ] }, { "id": 2402948, "author_name": "Abed Merii", "author_url": "", "post_date": "2023-08-22T12:16:36.810000", "content": "@readoc Thank you so much for sharing, most of the time we can learn more from the things that didn't work than from the things that we know will work.
\n.1. Do you have a link to the repo housing this code, or the notebook that you did this in? Would be curious to take a look
.4. What were your findings based on experimenting with the different loss functions?
", "votes": 1, "replies": [ { "id": 2404067, "author_name": "Sarthak Bhatt", "author_url": "", "post_date": "2023-08-23T03:26:43.217000", "content": "\"Making this notebook public will be locked until the competition ends\"
\nI am getting this message when trying to make the notebook public
Yeah, unfortunately I got the same error with my notebook :/
", "votes": 1, "replies": [] } ] } ] }, { "id": 2402519, "author_name": "hoho", "author_url": "", "post_date": "2023-08-22T07:33:45.707000", "content": "Wow! You rewrote ESM2 in JAX?! If that can save appreciable time or GPU memory, I think it will be very helpful! Did you have a benchmark result vs. pytorch-ESM2?
", "votes": 1, "replies": [ { "id": 2404068, "author_name": "", "author_url": "", "post_date": "2023-08-23T03:27:29.300000", "content": "", "votes": 1, "replies": [] } ] }, { "id": 2402477, "author_name": "IMP_imp", "author_url": "", "post_date": "2023-08-22T07:02:57.993000", "content": "Thanks for sharing, but I can't agree with the fourth point. I have similarly tried the fourth point you mentioned and this way can enhance the CV and LB scores and the enhancement is greater in BPO.
", "votes": 1, "replies": [] } ], "raw_markdown_by_id": { "2402415": "This competition was also the focus of my BTech Major Project and I have invested approximately 75 hours into this competition. The goal of the CAFA 5 competition is to predict biological function of proteins given their amino-acid sequences. One interesting part about the competition is that models are evaluated on \"future data,\" meaning the test set does not exist at submission and the ground truth is compiled months later from newly published lab results.\n\nI had two main strategies for winning this competition: \n1. **TPU-Maxxing**: Rewriting the ESM2-15B model in JAX to utilize TPUv3-8 pods via GSPMD tensor partitioning.\n2. **GNNs**: Using the Gene Ontology (GO) Hierarchy through Geometric Deep Learning techniques to improve embeddings for sparse \"tail\" classes.\n\n---\n\n## Scaling up to ESM-15B with TPUs\n\n[Kaggle Notebook: CAFA ESM JAX 15B](https://www.kaggle.com/code/readoc/cafa-esm-jax-15b/edit/run/133312976)\n\nIn Kaggle you can get either 30 hours of the P100 GPU or 20 hours of 8xTPUv3 for running modelling experiments. Simple math says that 20 hours of 8 TPUv3 will give you 13.2x more FLOPs than 30 hours of P100 GPUs. So there is a huge competitive advantage in being able to use TPUs effectively as compared to using consumer grade GPUs.\n\nRight now [ESM2](https://www.biorxiv.org/content/10.1101/2022.07.20.500902v2.full.pdf) (15B params) is the largest and the best protein language model available. However, it does not fit in a single TPU or a GPU so I rewrote ESM2 in JAX with [GSPMD](https://arxiv.org/pdf/2105.04663)-style tensor partitioning to fit the model across the 8 TPUs.\n\n**Embedding Strategy**: The main goal is to get the best possible embedding for a protein. I tried the last layer CLS embedding and also tried mean pooling to get 5k dim vector for every protein. Training an MLP on these embeddings gave me a score of ~0.54.\n\nOne failed experiment was that I tried to train a secondary transformer on mean embeddings per amino acid, but this did not outperform the baseline MLP.\n\n---\n\n## GO Hierarchy Aggregation\n\n[Kaggle Notebook (Very Rough State)](https://www.kaggle.com/code/readoc/cafa-keras-mlp-gnn-v2)\n\n**Tail Class Problem**: There are approximately 40k unique GO terms. Many specific functions (leaf nodes) have very few training samples. My hypothesis was that we could enrich these \"tail\" terms by aggregating information from their more frequent \"parent\" and \"sibling\" nodes.\n\nMy second secret strategy was using GNNs to build better protein embeddings by using the geometric structural dependencies within the Gene Ontology (GO) hierarchy. \n\n### Example Relationship between Proteins\n\n* **Protein A (Child Node)**: Annotated with a very specific rare leaf node. It is a tail class with low sample frequency.\n* **Protein B (Parent Node)**: Represents a high-level biological process. It has a high sample frequency in the dataset.\n* **Protein C (Sibling Node to Protein A)**: Another specialized sub-process sharing the same parent (Protein B) as Protein A.\n\nNow instead of treating every protein as independent and just using the amino acid sequence of the protein to build their embedding, we can now also use these parent, child, and sibling relationships between the proteins to build even better protein embeddings. \n\nI did not see anyone in this competition using this Gene Ontology structure so I took this as another competitive advantage. Using GNNs was also one of my main reasons for entering this competition since I have gotten very interested in Graph Neural Networks recently.\n\n### Neighborhood Aggregation Logic\n\nTo do this I focused on the top 3k term vectors and built a 3000x3000 adjacency matrix. Then I used power iteration to build a transition matrix:\n* **Child Influence**: 0.75 weight (Child -> Parent)\n* **Parent Influence**: 0.25 weight (Parent -> Child)\n\nMy code here was brute force but thanks to the power of TPUs it ended up working anyways. I also tried spectral position embedding but it did not give a better score. \n\n---\n\n## Other Tricks\n\n**Full Finetuning of the ProtT5 model**: I fully finetuned the ProtT5 model on TPUs. This also did not give a better score as compared to just using fixed embedding and training an MLP on it.\n\n**Different Loss Functions**: To address label imbalance, I tried different loss functions including focal loss, dice loss, poly loss, and the classic approach of up-weighting low-sample terms. All of these losses gave very similar scores which was a little surprising as I thought this would play a bigger role.", "2402552": "How can you tell than ESM 15B didn't work based on a score of 0.54? That seems like a strong public LB for a single model.", "2403867": "Surprised that your approach #3 didn't work better.\n\nIn my hands focal loss worked better than binary cross-entropy.", "2402948": "@readoc Thank you so much for sharing, most of the time we can learn more from the things that didn't work than from the things that we know will work.\n\n.1. Do you have a link to the repo housing this code, or the notebook that you did this in? Would be curious to take a look