Graph Set Transformer

2026-06-03

AI Summary

Existing graph set learning methods suffer from an information bottleneck due to the separation between graph embedding and set-level contextual modeling. This work proposes a novel end-to-end trainable architecture that interleaves node feature propagation via graph neural networks (GNNs) with cross-graph contextual modeling using Set Transformers at every layer, enabling hierarchical interaction between local structural and global set-level information. A gating mechanism is introduced to dynamically fuse these two sources of information, achieving their integration directly at the architectural level for the first time. The proposed method consistently outperforms baseline approaches across a synthetic task and three real-world datasets—reaction center identification, reaction yield prediction, and image classification. Ablation studies confirm that the intertwined design of local and set-level contextual modeling is crucial to the observed performance gains.

Abstract

We introduce the Graph Set Transformer (GST), a neural network architecture for learning on sets of graphs, designed for tasks in which per-element predictions depend on set-wide context as well as local structure. Existing architectures, including DeepSets and SetTransformer, require pre-encoded graph embeddings from a separate GNN, creating a bottleneck between feature extraction and set-level contextualisation. In contrast, GST interleaves node-level feature propagation and cross-graph contextual modelling at every layer, fusing the two levels of information through a gating mechanism. We evaluate GST on a controlled synthetic suite designed to isolate set-conditional structural reasoning and on three real-data benchmarks spanning per-atom reaction-centre identification, reaction yield prediction, and image classification. Under matched parameter budgets, GST performs better than the baselines across these settings. An architectural ablation strongly suggests that the interleaving of local and set context contributes substantially to this advantage.

Problem

• graph sets
• set-level context
• graph neural networks
• feature extraction bottleneck
• per-element prediction

Innovation

• Graph Set Transformer
• set-level contextualisation
• interleaved feature propagation
• gating mechanism
• graph representation learning