COMBA: Cross Batch Aggregation for Learning Large Graphs with Context Gating State Space Models

TL;DR

COMBA introduces a novel method combining graph context gating and cross batch aggregation to enable efficient learning on large graphs using state space models, outperforming existing approaches.

Contribution

The paper presents COMBA, a new approach that adapts state space models for large graph learning through innovative context gating and cross batch aggregation techniques.

Findings

Significant performance improvements over baseline methods.

Theoretical proof of lower error with cross-batch aggregation.

Scalability to large graphs demonstrated on benchmark datasets.

Abstract

State space models (SSMs) have recently emerged for modeling long-range dependency in sequence data, with much simplified computational costs than modern alternatives, such as transformers. Advancing SMMs to graph structured data, especially for large graphs, is a significant challenge because SSMs are sequence models and the shear graph volumes make it very expensive to convert graphs as sequences for effective learning. In this paper, we propose COMBA to tackle large graph learning using state space models, with two key innovations: graph context gating and cross batch aggregation. Graph context refers to different hops of neighborhood for each node, and graph context gating allows COMBA to use such context to learn best control of neighbor aggregation. For each graph context, COMBA samples nodes as batches, and train a graph neural network (GNN), with information being aggregated cross batches, allowing COMBA to scale to large graphs. Our theoretical study asserts that cross-batch aggregation guarantees lower error than training GNN without aggregation. Experiments on benchmark networks demonstrate significant performance gains compared to baseline approaches. Code and benchmark datasets will be released for public access.