| --- |
| license: mit |
| pipeline_tag: graph-ml |
| tags: |
| - chemistry |
| - biology |
| --- |
| |
| # ProteinBindingGNN (44M Parameters) |
|
|
| **ProteinBindingGNN** is a 44-million parameter hybrid Graph Neural Network (GNN) and Transformer architecture designed for node-level protein binding site prediction. It processes both the 3D geometric structure of a protein and its biochemical sequence context (via ESM-2 embeddings) to predict the probability of each amino acid residue belonging to a binding pocket. |
|
|
| The model combines the local geometric awareness of **Equivariant Graph Neural Networks (EGNNs)** with the global, pairwise reasoning capabilities of **Evoformer blocks** (inspired by AlphaFold 2), executed through a multi-step recycling mechanism. |
|
|
| ## Usage |
| You can find everything related to this model at [GraphBind](https://github.com/Labrapuerta/GraphBind) |
|
|
| ## 🧬 Model Architecture |
|
|
| The architecture is explicitly designed to balance rich feature preservation (high hidden dimensionality) with controlled reasoning depth to prevent over-smoothing on 3D protein graphs. |
|
|
| * **Parameters:** ~44.1 Million |
| * **Local Geometry:** 3 EGNN Layers (prevents message-passing over-smoothing, strictly capturing the 3-hop micro-environment). |
| * **Global Reasoning:** 8 Evoformer Blocks with Coordinate-Aware Edge Updates. |
| * **Recycling:** 5 iterations (the representations and coordinates are iteratively refined). |
|
|
| ### Full Configuration |
| ```python |
| CONFIG = { |
| "node_input_dim": 1305, # Designed for pre-trained ESM-2 embeddings |
| "edge_input_dim": 4, # e.g., distances, bond types, angles |
| "hidden_dim": 512, # Preserves rich biochemical context |
| "num_egnn_layers": 3, # Local 3D message passing |
| "num_evoformer_blocks": 8, # Global pairwise attention |
| "num_heads": 16, # Attention heads (d_k = 32) |
| "dropout": 0.3, # High regularization for generalization |
| "update_coords": True, # Dynamic coordinate refinement |
| "num_recycles": 5, # AlphaFold 2 style recycling loop |
| "alpha": 0.3, # Edge update mixing coefficient |
| } |
| ``` |
|
|
| ## Performance & Training Details |
| This model was trained on a highly curated dataset of 6,700 protein complexes. |
| You can see the full run and details at [Weights and Biases](https://wandb.ai/a01700257/Graph_binding/workspace?nw=nwusera01700257) |
|
|
| ### Class Imbalance Handling |
| Protein binding datasets suffer from extreme class imbalance (typically ~5% of residues are true binders, 95% are non-binders). The model is optimized for high ranking capability, achieving: |
|
|
| - Validation AUROC: 0.8898 |
|
|
| - Validation F1 (Standard Threshold): 0.3875 (Note: F1 is highly sensitive to the massive negative class; threshold tuning is recommended for downstream applications). |
|
|
| - Validation Recall: ~0.50 - 0.57 |
|
|
|
|
|
|
| ## Future work |
| Due to the limitations on gpu's, the model wans't trained enough to show the expected results, in the future with a sponsorship it is planned to keep training and fine tunning. |
