Models and Architectures

Overview

TorchDrug provides a comprehensive collection of pre-built model architectures for various graph-based learning tasks. This reference catalogs all available models with their characteristics, use cases, and implementation details.

Graph Neural Networks

GCN (Graph Convolutional Network)

Type: Spatial message passing Paper: Semi-Supervised Classification with Graph Convolutional Networks (Kipf & Welling, 2017)

Characteristics: - Simple and efficient aggregation - Normalized adjacency matrix convolution - Works well for homophilic graphs - Good baseline for many tasks

Best For: - Initial experiments and baselines - When computational efficiency is important - Graphs with clear local structure

Parameters: - input_dim: Node feature dimension - hidden_dims: List of hidden layer dimensions - edge_input_dim: Edge feature dimension (optional) - batch_norm: Apply batch normalization - activation: Activation function (relu, elu, etc.) - dropout: Dropout rate

Use Cases: - Molecular property prediction - Citation network classification - Social network analysis

GAT (Graph Attention Network)

Type: Attention-based message passing Paper: Graph Attention Networks (Veličković et al., 2018)

Characteristics: - Learns attention weights for neighbors - Different importance for different neighbors - Multi-head attention for robustness - Handles varying node degrees naturally

Best For: - When neighbor importance varies - Heterogeneous graphs - Interpretable predictions

Parameters: - input_dim, hidden_dims: Standard dimensions - num_heads: Number of attention heads - negative_slope: LeakyReLU slope - concat: Concatenate or average multi-head outputs

Use Cases: - Protein-protein interaction prediction - Molecule generation with attention to reactive sites - Knowledge graph reasoning with relation importance

GIN (Graph Isomorphism Network)

Type: Maximally powerful message passing Paper: How Powerful are Graph Neural Networks? (Xu et al., 2019)

Characteristics: - Theoretically most expressive GNN architecture - Injective aggregation function - Can distinguish graph structures GCN cannot - Often best performance on molecular tasks

Best For: - Molecular property prediction (state-of-the-art) - Tasks requiring structural discrimination - Graph classification

Parameters: - input_dim, hidden_dims: Standard dimensions - edge_input_dim: Include edge features - batch_norm: Typically use true - readout: Graph pooling ("sum", "mean", "max") - eps: Learnable or fixed epsilon

Use Cases: - Drug property prediction (BBBP, HIV, etc.) - Molecular generation - Reaction prediction

RGCN (Relational Graph Convolutional Network)

Type: Multi-relational message passing Paper: Modeling Relational Data with Graph Convolutional Networks (Schlichtkrull et al., 2018)

Characteristics: - Handles multiple edge/relation types - Relation-specific weight matrices - Basis decomposition for parameter efficiency - Essential for knowledge graphs

Best For: - Knowledge graph reasoning - Heterogeneous molecular graphs - Multi-relational data

Parameters: - num_relation: Number of relation types - hidden_dims: Layer dimensions - num_bases: Basis decomposition (reduce parameters)

Use Cases: - Knowledge graph completion - Retrosynthesis (different bond types) - Protein interaction networks

MPNN (Message Passing Neural Network)

Type: General message passing framework Paper: Neural Message Passing for Quantum Chemistry (Gilmer et al., 2017)

Characteristics: - Flexible message and update functions - Edge features in message computation - GRU updates for node hidden states - Set2Set readout for graph representation

Best For: - Quantum chemistry predictions - Tasks with important edge information - When node states evolve over multiple iterations

Parameters: - input_dim, hidden_dim: Feature dimensions - edge_input_dim: Edge feature dimension - num_layer: Message passing iterations - num_mlp_layer: MLP layers in message function

Use Cases: - QM9 quantum property prediction - Molecular dynamics - 3D conformation-aware tasks

SchNet (Continuous-Filter Convolutional Network)

Type: 3D geometry-aware convolution Paper: SchNet: A continuous-filter convolutional neural network (Schütt et al., 2017)

Characteristics: - Operates on 3D atomic coordinates - Continuous filter convolutions - Rotation and translation invariant - Excellent for quantum chemistry

Best For: - 3D molecular structure tasks - Quantum property prediction - Protein structure analysis - Energy and force prediction

Parameters: - input_dim: Atom features - hidden_dims: Layer dimensions - num_gaussian: RBF basis functions for distances - cutoff: Interaction cutoff distance

Use Cases: - QM9 property prediction - Molecular dynamics simulations - Protein-ligand binding with structures - Crystal property prediction

ChebNet (Chebyshev Spectral CNN)

Type: Spectral convolution Paper: Convolutional Neural Networks on Graphs (Defferrard et al., 2016)

Characteristics: - Spectral graph convolution - Chebyshev polynomial approximation - Captures global graph structure - Computationally efficient

Best For: - Tasks requiring global information - When graph Laplacian is informative - Theoretical analysis

Parameters: - input_dim, hidden_dims: Dimensions - num_cheb: Order of Chebyshev polynomial

Use Cases: - Citation network classification - Brain network analysis - Signal processing on graphs

NFP (Neural Fingerprint)

Type: Molecular fingerprint learning Paper: Convolutional Networks on Graphs for Learning Molecular Fingerprints (Duvenaud et al., 2015)

Characteristics: - Learns differentiable molecular fingerprints - Alternative to hand-crafted fingerprints (ECFP) - Circular convolutions like ECFP - Interpretable learned features

Best For: - Molecular similarity learning - Property prediction with limited data - When interpretability is important

Parameters: - input_dim, output_dim: Feature dimensions - hidden_dims: Layer dimensions - num_layer: Circular convolution depth

Use Cases: - Virtual screening - Molecular similarity search - QSAR modeling

Protein-Specific Models

GearNet (Geometry-Aware Relational Graph Network)

Type: Protein structure encoder Paper: Protein Representation Learning by Geometric Structure Pretraining (Zhang et al., 2023)

Characteristics: - Incorporates 3D geometric information - Multiple edge types (sequential, spatial, KNN) - Designed specifically for proteins - State-of-the-art on protein tasks

Best For: - Protein structure prediction - Protein function prediction - Protein-protein interaction - Any task with protein 3D structures

Parameters: - input_dim: Residue features - hidden_dims: Layer dimensions - num_relation: Edge types (sequence, radius, KNN) - edge_input_dim: Geometric features (distances, angles) - batch_norm: Typically true

Use Cases: - Enzyme function prediction (EnzymeCommission) - Protein fold recognition - Contact prediction - Binding site identification

ESM (Evolutionary Scale Modeling)

Type: Protein language model (transformer) Paper: Biological structure and function emerge from scaling unsupervised learning (Rives et al., 2021)

Characteristics: - Pre-trained on 250M+ protein sequences - Captures evolutionary and structural information - Transformer architecture - Transfer learning for downstream tasks

Best For: - Any sequence-based protein task - When no structure available - Transfer learning with limited data

Variants: - ESM-1b: 650M parameters - ESM-2: Multiple sizes (8M to 15B parameters)

Use Cases: - Protein function prediction - Variant effect prediction - Protein design - Structure prediction (ESMFold)

ProteinBERT

Type: Masked language model for proteins

Characteristics: - BERT-style pre-training - Masked amino acid prediction - Bidirectional context - Good for sequence-based tasks

Use Cases: - Function annotation - Subcellular localization - Stability prediction

ProteinCNN / ProteinResNet

Type: Convolutional networks for sequences

Characteristics: - 1D convolutions on sequences - Local pattern recognition - Faster than transformers - Good for motif detection

Use Cases: - Binding site prediction - Secondary structure prediction - Domain identification

ProteinLSTM

Type: Recurrent network for sequences

Characteristics: - Bidirectional LSTM - Captures long-range dependencies - Sequential processing - Good baseline for sequence tasks

Use Cases: - Order prediction - Sequential annotation - Time-series protein data

Knowledge Graph Models

TransE (Translation Embedding)

Type: Translation-based embedding Paper: Translating Embeddings for Modeling Multi-relational Data (Bordes et al., 2013)

Characteristics: - h + r ≈ t (head + relation ≈ tail) - Simple and interpretable - Works well for 1-to-1 relations - Memory efficient

Best For: - Large knowledge graphs - Initial experiments - Interpretable embeddings

Parameters: - num_entity, num_relation: Graph size - embedding_dim: Embedding dimensions (typically 50-500)

RotatE (Rotation Embedding)

Type: Rotation in complex space Paper: RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space (Sun et al., 2019)

Characteristics: - Relations as rotations in complex space - Handles symmetric, antisymmetric, inverse, composition - State-of-the-art on many benchmarks

Best For: - Most knowledge graph tasks - Complex relation patterns - When accuracy is critical

Parameters: - num_entity, num_relation: Graph size - embedding_dim: Must be even (complex embeddings) - max_score: Score clipping value

DistMult

Type: Bilinear model

Characteristics: - Symmetric relation modeling - Fast and efficient - Cannot model antisymmetric relations

Best For: - Symmetric relations (e.g., "similar to") - When speed is critical - Large-scale graphs

ComplEx

Type: Complex-valued embeddings

Characteristics: - Handles asymmetric and symmetric relations - Better than DistMult for most graphs - Good balance of expressiveness and efficiency

Best For: - General knowledge graph completion - Mixed relation types - When RotatE is too complex

SimplE

Type: Enhanced embedding model

Characteristics: - Two embeddings per entity (canonical + inverse) - Fully expressive - Slightly more parameters than ComplEx

Best For: - When full expressiveness needed - Inverse relations are important

Generative Models

GraphAutoregressiveFlow

Type: Normalizing flow for molecules

Characteristics: - Exact likelihood computation - Invertible transformations - Stable training (no adversarial) - Conditional generation support

Best For: - Molecular generation - Density estimation - Interpolation between molecules

Parameters: - input_dim: Atom features - hidden_dims: Coupling layers - num_flow: Number of flow transformations

Use Cases: - De novo drug design - Chemical space exploration - Property-targeted generation

Pre-training Models

InfoGraph

Type: Contrastive learning

Characteristics: - Maximizes mutual information - Graph-level and node-level contrast - Unsupervised pre-training - Good for small datasets

Use Cases: - Pre-train molecular encoders - Few-shot learning - Transfer learning

MultiviewContrast

Type: Multi-view contrastive learning for proteins

Characteristics: - Contrasts different views of proteins - Geometric pre-training - Uses 3D structure information - Excellent for protein models

Use Cases: - Pre-train GearNet on protein structures - Transfer to property prediction - Limited labeled data scenarios

Model Selection Guide

By Task Type

Molecular Property Prediction: 1. GIN (first choice) 2. GAT (interpretability) 3. SchNet (3D available)

Protein Tasks: 1. ESM (sequence only) 2. GearNet (structure available) 3. ProteinBERT (sequence, lighter than ESM)

Knowledge Graphs: 1. RotatE (best performance) 2. ComplEx (good balance) 3. TransE (large graphs, efficiency)

Molecular Generation: 1. GraphAutoregressiveFlow (exact likelihood) 2. GCPN with GIN backbone (property optimization)

Retrosynthesis: 1. GIN (synthon completion) 2. RGCN (center identification with bond types)

By Dataset Size

Small (< 1k): - Use pre-trained models (ESM for proteins) - Simpler architectures (GCN, ProteinCNN) - Heavy regularization

Medium (1k-100k): - GIN for molecules - GAT for interpretability - Standard training

Large (> 100k): - Any model works - Deeper architectures - Can train from scratch

By Computational Budget

Low: - GCN (simplest) - DistMult (KG) - ProteinLSTM

Medium: - GIN - GAT - ComplEx

High: - ESM (large) - SchNet (3D) - RotatE with high dim

Implementation Tips

1. Start Simple: Begin with GCN or GIN baseline
2. Use Pre-trained: ESM for proteins, InfoGraph for molecules
3. Tune Depth: 3-5 layers typically sufficient
4. Batch Normalization: Usually helps (except KG embeddings)
5. Residual Connections: Important for deep networks
6. Readout Function: "mean" usually works well
7. Edge Features: Include when available (bonds, distances)
8. Regularization: Dropout, weight decay, early stopping