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

Jiajun Shen; Yufei Jin; Yi He; xingquan Zhu

cs.LGFeb 19, 2026

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

Authors: Jiajun Shen, Yufei Jin, Yi He, xingquan Zhu

Organizations: Florida Atlantic University · William & Mary

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.

Figures & tables

Figure 1: Cross batch aggregation process. From left to right, given a graph in ①, the nodes are partitioned into batches (only two batches are shown in ②). For each batch, COMBA first finds each nodes’

\hat{k}

-hop neighbors and forms a subgraph shown in ③ (

\hat{k}=2

in this case). A GNN is trained for each 1-hope, 2-hop, and

k

-hop based adjacency matrix. When training each hop’s GNN, information from other batches are used to help learn current batch node’s embedding. E.g. , in ④, node

a_{2}

in Batch 1 aggregates information from

a_{4}

from Batch 2. Cross batch aggregation allows all GNNs being trained to collectively help each others.

Figure 2: Illustration of the COMBA block with context gating. From left to right. The input sequence

\mathcal{Z^{\prime}}

is first processed by the S4 module to produce hop-wise representations. A context gating mechanism C is then applied over the local hop window

z_{k-w:k+w}

to refine each hop embedding

y_{k}

. The gated outputs are concatenated with the original node features, forming a new sequence

Y

. Pooling along the hop dimension aggregates the sequence into embedding

X^{\prime}

for downstream tasks.

Figure 3: The proposed COMBA framework on large homogeneous graph. From left to right. Nodes of a homogeneous graph in ① are partitioned into

\hat{m}

batches in ②. ③: for each batch, COMBA identifies the

\hat{k}

-hop neighbors of target nodes and constructs a corresponding subgraph. ④: Node embeddings will be updated across batches via cross batch aggregation as illustrated in Fig. 1 . The resulting sequence will pass into COMBA block with context gating ⑤ as illustrated in Fig. 2 . ⑥: the final predictions

\hat{Y}

for all nodes are obtained and optimized using the cross-entropy loss

Figure 4: Log average runtime per epoch ( $y$ -axis) using fixed batch sizes vs. the sum of number of nodes and number of edges in log scale ( $x$ -axis).

Model	Roman-empire Accuracy	Amazon-ratings Accuracy	Ogbn-arxiv Accuracy	Ogbn-product Accuracy	Minesweeper ROC AUC	Tolokers ROC AUC
GCN	0.795 ${}_{\pm 0.0088}^{*}$	0.462 ±0.0080	0.680 ${}_{\pm 0.0044}^{*}$	OOM	0.884 ${}_{\pm 0.0046}^{*}$	0.838 ${}_{\pm 0.0034}^{*}$
Gated-GCN	0.833 ${}_{\pm 0.0079}^{*}$	0.481 ${}_{\pm 0.0071}^{*}$	0.701 ${}_{\pm 0.0065}^{*}$	OOM	0.905 ${}_{\pm 0.0016}^{*}$	0.832 ${}_{\pm 0.0090}^{*}$
Nagphormer	0.800 ${}_{\pm 0.0043}^{*}$	0.504 ±0.0070	0.696 ${}_{\pm 0.0023}^{*}$	OOM	0.903 ${}_{\pm 0.0006}^{*}$	0.834 ${}_{\pm 0.0045}^{*}$
Cluster-GCN	0.815 ${}_{\pm 0.0042}^{*}$	0.476 ${}_{\pm 0.0047}^{*}$	0.678 ${}_{\pm 0.0076}^{*}$	0.717 ${}_{\pm 0.0053}^{*}$	0.886 ${}_{\pm 0.0024}^{*}$	0.758 ${}_{\pm 0.0056}^{*}$
Graph Mamba-1	0.677 ${}_{\pm 0.0009}^{*}$	0.415 ${}_{\pm 0.0016}^{*}$	0.604 ${}_{\pm 0.0020}^{*}$	OOM	0.806 ${}_{\pm 0.0059}^{*}$	0.734 ${}_{\pm 0.0113}^{*}$
Graph Mamba-2	0.869 ${}_{\pm 0.0092}^{*}$	0.490 ${}_{\pm 0.0036}^{*}$	0.686 ${}_{\pm 0.0089}^{*}$	OOM	0.927 ${}_{\pm 0.0021}^{*}$	0.803 ${}_{\pm 0.0199}^{*}$

Table 1: Performance comparisons between baselines and our proposed method across five homogeneous datasets. Over 5 different initialization status, accuracies are reported for Roman-empire, Amazon-ratings, Ogbn-arxiv, and Ogbn-products and roc auc scores are reported for Minesweeper and Tolokers. Superscript * indicates that COMBA is statistically significantly better than this method at 95% confidence level using the performance metrics.

Model	Roman-empire Accuracy	Amazon-ratings Accuracy	Ogbn-arxiv Accuracy	Minesweeper ROC AUC	Tolokers ROC AUC
COMBA	0.895 ±0.0038	0.507 ±0.0025	0.716 ±0.0037	0.942 ±0.0039	0.845 ±0.0009
w/o cross batch	0.881 ±0.0040	0.492 ±0.0062	0.703 ±0.0009	0.940 ±0.0051	0.832 ±0.0022
w/o context gating	0.878 ±0.0005	0.472 ±0.0007	0.705 ±0.0014	0.939 ±0.0050	0.805 ±0.0039

Table 2: Ablation study results w.r.t. cross batch aggregation and context gating mechanism

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

	Roman-empire	Amazon-ratings	Minesweeper	Tolokers	Ogbn-arxiv	Ogbn-products
# Nodes	22,662	24,492	10,000	11,758	169,343	2,449,029
# Edges	32,927	93,050	39,402	51,900	1,166,243	61,859,140
# Features	300	300	7	10	128	100
# Classes	18	5	2	2	40	47

Appendix

Table 3: Dataset Statistics

Explore similar work

Learning Long Range Spatio-Temporal Representations over Continuous Time Dynamic Graphs with State Space Models
Jun 3, 2026Ayushman Raghuvanshi, Thummaluru Siddartha Reddy, Sundeep Prabhakar Chepuri +1Long-Context Modeling Temporal Link Prediction
Department of Electrical Communication Engineering, Indian Institute of Science, Bangalore
Graph Signal Processing Meets Mamba2: Adaptive Filter Bank via Delta Modulation
Mar 20, 2026Yehjin Shin, Seojin Kim, Noseong ParkSelective SSMs Long-Context Language Modeling
KAIST
Are Large Language Models Suitable for Graph Computation? Progress and Prospects
Jun 5, 2026Yuting Zhang, Yi Han, Kai Wang +3LLM Reasoning with Graphs LLM Tool Use
University of New South Wales · Antai College of Economics and Management, Shanghai Jiao Tong University · Edith Cowan University +1

Jun 3, 2026J/K move · Enter open · S save

Learning Long Range Spatio-Temporal Representations over Continuous Time Dynamic Graphs with State Space Models

Ayushman Raghuvanshi, Thummaluru Siddartha Reddy, Sundeep Prabhakar Chepuri, Mahesh Chandran

Department of Electrical Communication Engineering, Indian Institute of Science, Bangalore

Long-Context Modeling Temporal Link Prediction Dynamic Graph Learning State Space Models

Open paper AskPDF arXiv

Continuous-time dynamic graphs (CTDGs) provide a richer framework to capture fine-grained temporal patterns in evolving relational data. Long-range information propagation is a key challenge while learning representations, wherein it is important to retain and update information over long temporal horizons. Existing approaches restrict models to capture one-hop or local temporal neighborhoods and fail to capture multi-hop or global structural patterns. To mitigate this, we derive a parameter-efficient state-space modeling framework for continuous-time dynamic graphs (CTDG-SSM) from first principles. We first introduce continuous-time Topology-Aware higher order polynomial projection operator (CTT-HiPPO), a novel memory-based reformulation of HiPPO to jointly encode temporal dynamics and graph structure. The solution from CTT-HiPPO is obtained by projecting the classical HiPPO solution through a polynomial of the Laplacian matrix, yielding topology-aware memory updates that admit an equivalent state-space formulation for CTDGs (CTDG-SSM). Then a computationally efficient discrete formulation is obtained using the zero-order hold approach for model implementation. Across benchmarks on dynamic link prediction, dynamic node classification, and sequence classification, CTDG-SSM achieves state-of-the-art performance. Notably, it achieves large performance gains on datasets that require long range temporal (LRT) and spatial reasoning.

Abstract

Explore similar work

Learning Long Range Spatio-Temporal Representations over Continuous Time Dynamic Graphs with State Space Models

Graph Signal Processing Meets Mamba2: Adaptive Filter Bank via Delta Modulation

Are Large Language Models Suitable for Graph Computation? Progress and Prospects