Graph Set Transformer: Learning Graph Representations with Set Context
Abstract
We introduce the Graph Set Transformer (GST), a neural network architecture for learning on sets of graphs, motivated by the assumption that structurally heterogeneous graphs may share higher-level semantics. Existing architectures, including DeepSets and Set Transformer, require pre-encoded graph embeddings from a separate GNN, creating a bottleneck between feature extraction and set-level contextualisation. In contrast, GST avoids this bottleneck by performing node-level feature propagation and cross-graph contextual modelling simultaneously, fusing the two levels of information through a gating mechanism. Across five molecular classification benchmarks, GST achieves consistent ROC-AUC scores of 98.5-99.6% compared to the best baselines of 89.3-98.8% for large sets of cardinalities 10 and 20. In a drug-drug interaction benchmark on sparse, undersampled data, GST improves the F1 score by 17.5% compared to the best baseline.