Module 3: Heterogeneous Graph Transformer (HGT)

With the modalities perfectly aligned via Module 2, ProMaya must encode this dense information into a single, comprehensive representation per protein. Standard Graph Neural Networks (GNNs) assume all nodes are of the same type. ProMaya utilizes a Heterogeneous Graph Transformer (HGT) to handle the disparate modalities simultaneously.

Graph Topology

ProMaya constructs a massive heterogeneous graph (\(G\)) for the protein containing:

4 Node Types

Atom: Projected from the 41-d atomic graph.
Residue: Projected from the 1082-d residue features.
Surface: Projected from the 128-d PointNet++ encodings.
Sequence: Projected from the 1024-d ProtTrans embeddings.

Note: All node features are linearly projected to a uniform 256-d dimension before entering the HGT.

6 Edge Relations

The network defines specific message-passing pathways between node types: - atom-atom (covalent and non-covalent proximity) - residue-residue (\(C_\alpha\) proximity) - atom-residue (hierarchical containment) - residue-surface (spatial projection) - sequence-residue (1D to 3D mapping) - cross-level (auxiliary alignment edges)

HGT Architecture

The HGT encoder applies 4 layers of message passing. Unlike standard Transformers, the Query, Key, and Value matrices are type-specific.

For a target node \(t\) and source node \(s\) connected by edge \(e\): 1. Attention Calculation: Uses a type-specific projection matrix \(W^{ATT}_{\tau(s), \tau(t)}\) based on the source and target node types. 2. Message Passing: Uses a relation-specific weight matrix \(W^{MSG}_{\phi(e)}\). 3. Aggregation: Multi-head attention (8 heads) aggregates messages across the varying neighborhoods. 4. Update: Residual connections, LayerNorm, and a Feed-Forward Network (512-d, GELU) update the node embedding.

src/models/hgt.py

import torch
import torch.nn as nn
from torch_geometric.nn import HGTConv

class ProMayaHGTEncoder(nn.Module):
    def __init__(self, metadata, hidden_channels=256, out_channels=256, num_heads=8, num_layers=4):
        super().__init__()

        # Linear projection to ensure all modalities have dimension = 256 before HGT
        self.lin_dict = nn.ModuleDict({
            node_type: nn.Linear(in_dim, hidden_channels)
            for node_type, in_dim in metadata['node_dims'].items()
        })

        self.convs = nn.ModuleList()
        for _ in range(num_layers):
            conv = HGTConv(hidden_channels, hidden_channels, metadata['graph_metadata'],
                           num_heads, group='sum')
            self.convs.append(conv)

        self.out_lin = nn.Linear(hidden_channels, out_channels)

    def forward(self, x_dict, edge_index_dict):
        # 1. Project to uniform hidden dimension
        x_dict = {
            node_type: self.lin_dict[node_type](x) 
            for node_type, x in x_dict.items()
        }

        # 2. 4 Layers of Heterogeneous Message Passing
        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)

        # 3. Output embeddings
        return {
            node_type: self.out_lin(x) 
            for node_type, x in x_dict.items()
        }

Protein-Level Embedding

After 4 layers, the graph contains highly refined embeddings. A modality-specific attention pooling mechanism extracts the most salient features from all 4 node types, concatenates them (\(4 \times 256 = 1024\text{-d}\)), and linearly projects them down to a final 512-dimensional protein embedding (\(e_P\)).

Next: Move from single proteins to protein pairs in Cross-Protein Interaction â†’