Heterogeneous Graph Learning with TGraphX: RGCN, HAN, and HGT
Most graph learning tutorials treat graphs as homogeneous: every node is the same type, every edge is the same type. Many real-world graphs are heterogeneous — they have multiple node types (papers, authors, venues) and multiple edge types (writes, published_at, cites). Modeling type heterogeneity correctly can meaningfully improve performance on link prediction and node classification tasks.
TGraphX provides Relational GCN (RGCN), Heterogeneous Attention Network (HAN), and Heterogeneous Graph Transformer (HGT) layers, along with a HeteroGraph data structure. These are labeled Experimental — the APIs may change — but they are functional for research use.
Understanding Heterogeneous Graphs
A heterogeneous graph G = (V, E, T_v, T_e) has:
- Node types
T_v = {paper, author, venue, ...} - Edge types
T_e = {writes, cites, published_at, ...}
Each edge type connects specific source and destination node types. For example:
- (author, writes, paper)
- (paper, cites, paper)
- (paper, published_at, venue)
This type structure means different layers should handle different edge types separately and then combine their outputs.
TGraphX's HeteroGraph Data Structure
import torch
from tgraphx.core.hetero_graph import HeteroGraph
# Node features per type
node_features = {
"paper": torch.randn(100, 64), # 100 papers with 64-dim features
"author": torch.randn(50, 32), # 50 authors with 32-dim features
"venue": torch.randn(10, 16), # 10 venues with 16-dim features
}
# Edge indices per relation type
edge_index_dict = {
("author", "writes", "paper"): torch.randint(0, 50, (2, 200), dtype=torch.long),
("paper", "cites", "paper"): torch.randint(0, 100, (2, 500), dtype=torch.long),
("paper", "published_at", "venue"): torch.randint(0, 100, (2, 100), dtype=torch.long),
}
# Create heterogeneous graph
hg = HeteroGraph(
node_features=node_features,
edge_index_dict=edge_index_dict,
)
print(hg.node_types) # ['paper', 'author', 'venue']
print(hg.edge_types) # [('author', 'writes', 'paper'), ...]
print(hg.num_nodes("paper")) # 100
print(hg.num_edges("author", "writes", "paper")) # 200
Relational GCN (RGCN)
RGCN (Schlichtkrull et al., 2018) learns a separate weight matrix for each relation type and sums the transformed messages:
h_v^(k+1) = σ( Σ_{r ∈ R} Σ_{u ∈ N_r(v)} 1/|N_r(v)| W_r^(k) h_u^(k) + W_0^(k) h_v^(k) )
from tgraphx.layers.rgcn import RGCNConv
import torch.nn as nn
import torch.nn.functional as F
class RGCNEncoder(nn.Module):
def __init__(self):
super().__init__()
self.rgcn = RGCNConv(
in_dim=64,
out_dim=32,
num_relations=3, # writes, cites, published_at
num_bases=None, # None = full parameterization (no basis decomposition)
)
def forward(self, x, edge_index, edge_type):
return F.relu(self.rgcn(x, edge_index, edge_type))
# edge_type: [E] LongTensor with relation index for each edge
# Combine all edge types into one edge_index with type labels
all_edges = []
all_types = []
for rel_idx, (src_t, rel, dst_t) in enumerate(edge_index_dict.keys()):
ei = edge_index_dict[(src_t, rel, dst_t)]
all_edges.append(ei)
all_types.append(torch.full((ei.shape[1],), rel_idx, dtype=torch.long))
combined_ei = torch.cat(all_edges, dim=1) # [2, total_E]
combined_type = torch.cat(all_types) # [total_E]
# For RGCN on paper nodes
model = RGCNEncoder()
paper_x = node_features["paper"]
out = model(paper_x, combined_ei, combined_type)
print(out.shape) # [100, 32]
Heterogeneous Attention Network (HAN)
HAN (Wang et al., 2019) uses meta-path-based attention. A meta-path is a sequence of edge types that connects nodes of the same type through intermediate types. For example, Paper → Author → Paper (two papers written by the same author) is a meta-path.
from tgraphx.layers.han import HANConv
import torch.nn as nn
import torch.nn.functional as F
# Meta-path: paper-author-paper (papers co-authored by same author)
# This requires constructing the meta-path adjacency matrix from the raw edges
# TGraphX's HAN takes the meta-path edge index directly
# For simplicity: paper-author-paper adjacency (precomputed)
paper_via_author = torch.randint(0, 100, (2, 300), dtype=torch.long)
han = HANConv(
in_channels=64, # paper feature dim
out_channels=32,
num_heads=4,
)
paper_x = node_features["paper"]
out = han(paper_x, paper_via_author)
print(out.shape) # [100, 32]
Note: HAN requires constructing meta-path adjacency matrices before passing to the layer. TGraphX's HAN layer takes a pre-computed meta-path edge index, not raw heterogeneous edges. Meta-path construction must be done by the user.
Heterogeneous Graph Transformer (HGT)
HGT (Hu et al., 2020) uses type-specific linear transformations for keys, queries, and values, along with mutual attention across node types:
from tgraphx.layers.hgt import HGTConv
hgt = HGTConv(
in_channels={"paper": 64, "author": 32, "venue": 16},
out_channels=32,
metadata=(
list(node_features.keys()), # node types
list(edge_index_dict.keys()), # edge types (src_type, rel, dst_type)
),
heads=4,
)
# HGT processes all node types simultaneously
out_dict = hgt(node_features, edge_index_dict)
print(out_dict["paper"].shape) # [100, 32]
print(out_dict["author"].shape) # [50, 32]
Heterogeneous Graph Batching
For mini-batch training on heterogeneous graphs:
from tgraphx.core.hetero_batch import HeteroBatch
# Create a batch from multiple heterogeneous graphs
graphs = [hg, hg] # typically different graphs
batch = HeteroBatch.from_list(graphs)
print(batch.num_graphs) # 2
print(batch.num_nodes("paper")) # 200 (2 × 100)
Typed Neighbor Sampling
TGraphX includes typed neighbor sampling for heterogeneous graphs:
from tgraphx.hetero_sampling import HeteroNeighborSampler
sampler = HeteroNeighborSampler(
hg,
fanouts={"writes": 10, "cites": 5, "published_at": 3},
batch_size=16,
seed=42,
)
for batch in sampler:
# batch.node_features: dict of {node_type: Tensor}
# batch.edge_index_dict: dict of {(src, rel, dst): Tensor}
pass
Limitations and API Stability
Heterogeneous support is Experimental. The HeteroGraph, HGTConv, HANConv, and typed sampling APIs may change in future TGraphX releases. Do not build production systems on them without pinning the TGraphX version.
Meta-path construction is not automated. HAN requires pre-computed meta-path adjacency matrices. TGraphX does not include a meta-path extraction utility; this must be done externally.
HGT has higher memory cost. HGT's type-specific attention mechanism uses num_node_types × num_edge_types × heads × d_k attention parameters per layer. For graphs with many types, memory can be significant.
No built-in heterogeneous benchmark datasets. TGraphX does not include IMDB, DBLP, OAG, or other standard heterogeneous graph benchmarks. Load them via PyG or DGL and convert to HeteroGraph.
Different from DGL's HeteroGraph API. TGraphX's HeteroGraph is not API-compatible with DGL's HeteroGraph or PyG's HeteroData. Migration requires explicit conversion code.
When to Use TGraphX for Heterogeneous Graphs
TGraphX's heterogeneous support is most useful when:
- You already have a TGraphX workflow (tensor-valued node features, reproducibility tools, mining) and want to extend it to heterogeneous data.
- Your heterogeneous graph has tensor-valued node features (image entities, volumetric data) that require spatial message passing.
- You are prototyping with RGCN or HGT and do not need production-grade distributed hetero sampling.
For production heterogeneous GNN systems on large-scale data (OGBN-Mag, OAG), DGL's DistGraph with heterogeneous support or PyG's HeteroData with their heterogeneous samplers are more battle-tested choices.
Related Articles
- What is a TGX graph — the homogeneous data model
- Knowledge graph embedding with tensor features — typed entities in KGs
- TGraphX vs PyTorch Geometric — comparison of heterogeneous graph support
- Neighbor sampling with TGraphX — scalable sampling foundation
- Multimodal graph nodes in TGraphX — mixed modality node features