(Conceptual visualization of the sparse graph network)

A general framework for predicting protein variant activity from a small experimental dataset (tens of labelled samples). The method fuses pre-trained sequence embeddings (ESM2) with predicted structural pair embeddings (AlphaFold3) in a sparse edge-conditioned graph neural network.

Though developed and validated on two protein binder design cases, every component is protein-agnostic and can be applied to any activity-prediction task where per-residue single embeddings and residue-pair embeddings are available.

The core idea is to model a protein as a sparse graph in which:

• each node carries an ESM2 sequence embedding that encodes evolutionary and functional context,
• each edge carries an AF3 pair embedding that encodes the predicted spatial relationship between two residues,
• only the Top-K most informative edges (ranked by L2 norm of the pair embedding) are retained, removing noise and reducing computation.

An edge-conditioned GNN propagates information across this graph, an attention pooling layer collapses per-residue representations into a single graph-level vector, and a small MLP maps that vector to a predicted activity scalar.

ESM2 layer 32 (N × 1280)          AF3 pair embeddings (N × N × 128)
        │                                       │
   Linear projection                       Linear projection
   1280 → 256                              128 → 256
        │                                       │
        └───────────────┬───────────────────────┘
                        ↓
          ┌─────────────────────────────────┐
          │    EdgeConditionedGraphConv      │  × 2
          │                                 │
          │  h_i' = W_self · h_i            │
          │       + Σ_{j ∈ TopK(i)}         │
          │           α_ij · W_msg · [h_j ‖ e_ij] │
          │                                 │
          │  α_ij = softmax(MLP([h_i, h_j, e_ij])) │
          └─────────────────────────────────┘
                        ↓
               Attention Pooling
               g = Σ_i softmax(a_i) · h_i   →  (256,)
                        ↓
                MLP: 256 → 128 → 64 → 1
                        ↓
               Predicted activity (scalar)

Training uses 5-fold cross-validation with early stopping and is evaluated by Spearman ρ between predicted and measured activity values.

Despite the extremely small training sets, the method achieves meaningful Spearman correlations and substantially enriches high-activity sequences in the top-ranked candidates (e.g. 10.9× enrichment of high-activity NTD binders in the Top-500 predictions).

All numbers are Spearman ρ, 5-fold CV mean, averaged across multiple random seeds (8 seeds for NTD unless noted).

ESM2 node features are more informative and more stable than AF3 single embeddings. Concatenating both hurts in this low-data regime due to increased dimensionality.

Sparse graphs consistently outperform the full graph. Removing noisy low-importance edges is more beneficial than retaining all pairwise information. K=8 is the optimal operating point.

When the binding partner is available, running AF3 in complex mode yields pair embeddings that better reflect binding-relevant residue interactions.

The GNN message-passing step is essential — without it, pair embeddings cannot be exploited and performance collapses. The optional Scalar Mix across ESM2 layers yields a marginal improvement at 8× memory cost and is not recommended by default.

NTD_experiment/
├── README.md
├── NTD_activity.csv                      # Labelled training data (NTD case)
├── NTD_sequences.csv
│
├── Systematic_Experiments/               # Feature engineering baselines
│   ├── prepare_all_features.py
│   ├── run_all_experiments.py
│   └── features/
│       ├── esm2_single.csv               # (N × 1280)
│       ├── af3_single.csv                # (N × 384)
│       └── af3_pair.csv                  # (N × 256, mean+max pooled)
│
├── ESM2_Sequence_Models/                 # BiGRU / TCN sequence baselines
│   ├── models.py
│   └── train_sequence_models.py
│
├── AF3_GNN/                              # ⭐ Core GNN implementation
│   ├── gnn_models.py                     # Dense EdgeConditionedGNN (early version)
│   ├── optimized_gnn_models.py           # ⭐ Sparse GNN + ScalarMix (recommended)
│   ├── train_optimized_gnn.py            # 5-fold CV training
│   ├── train_with_receptor_5fold.py      # Complex-mode AF3 variant
│   ├── multi_seed_evaluation.py          # Multi-seed evaluation harness
│   ├── prepare_topk_graph_data.py        # Build sparse graph pickle files
│   ├── extract_esm2_layers.py            # Extract multi-layer ESM2 embeddings
│   └── data/
│       ├── single_layer_topk8_graph.pkl  # Default graph (36 MB)
│       ├── single_layer_full_graph.pkl   # Dense baseline (160 MB)
│       └── with_receptor_*.pkl           # Complex-mode variants
│
├── Candidates_Prediction/                # Inference on large candidate pools
│   ├── predict_with_5fold_ensemble.py
│   └── candidates_10874_predictions.csv
│
├── CHD_experiment/                       # Second validation case (CHD binders)
│
└── best_models/                          # Saved checkpoints
    └── {config}_{seed}/
        └── fold{1..5}_best.pth

conda activate ncaa   # PyTorch, ESM2, see ../ncaa_env.yml

1. Prepare node features — extract ESM2 layer-32 embeddings (shape N × 1280) for each protein variant.
2. Prepare edge features — run AlphaFold3 and extract pair embeddings (shape L × L × 128).
3. Build sparse graphs — use prepare_topk_graph_data.py as a template to assemble pickle files with keys node_features, edge_features, edge_index, and label.
4. Train and evaluate:

cd AF3_GNN
python multi_seed_evaluation.py --config single_layer_topk8 --num_seeds 8

cd AF3_GNN
# Default: ESM2 nodes + monomer AF3 pair + Top-K=8
python multi_seed_evaluation.py --config single_layer_topk8 --num_seeds 8

# With-receptor variant (best enrichment)
python train_with_receptor_5fold.py

cd Candidates_Prediction
python predict_with_5fold_ensemble.py
# Outputs: candidates_10874_predictions.csv  (ranked by predicted activity)

Each sample is stored as a Python dict in a pickle file:

{
    'node_features': np.ndarray,   # (seq_len, node_dim)
    'edge_features': np.ndarray,   # (num_edges, edge_dim)
    'edge_index':    np.ndarray,   # (2, num_edges)  — None for dense graph
    'label':         float,        # activity value
    'name':          str,          # variant identifier
}

Sparse > dense. Retaining only the K highest-L2-norm pair edges removes irrelevant long-range interactions and acts as a built-in regulariser, which is critical when training data is scarce.

Pre-trained features outperform task-specific learning. ESM2 and AF3 embeddings carry far more information than can be learned from dozens of labelled samples. The GNN learns to route and aggregate this information rather than learning representations from scratch.

Structure-aware pooling is essential. Replacing the GNN with global average pooling drops Spearman ρ by 76%, confirming that spatial message-passing is what makes pair embeddings useful.

Stability over peak performance. In low-data regimes, methods with low cross-seed variance (fixed L2-norm edge selection, no learnable edge scorer) are more reliable than methods that achieve higher peak performance on lucky seeds.

Training set size is the primary bottleneck. Both validated cases use fewer than 40 labelled samples, leading to a large train/validation performance gap (train ρ ≈ 0.95, val ρ ≈ 0.50). Collecting 200–500 labelled variants per task is expected to significantly improve generalisation. Until then, the Top-K sparse graph and the choice of strong pre-trained features are the most effective mitigation strategies available.

@software{sparse_ecgnn_2025,
  title  = {Sparse Edge-Conditioned GNN for Small-Data Protein Activity Prediction},
  year   = {2025},
  url    = {https://github.com/your-repo/sparse-ecgnn}
}