Table of Contents
Fetching ...

HeteGraph-Mamba: Heterogeneous Graph Learning via Selective State Space Model

Zhenyu Pan, Yoonsung Jeong, Xiaoda Liu, Han Liu

TL;DR

HeteGraph-Mamba (HGMN) addresses the challenge of learning from large heterogeneous graphs by marrying selective state space models with a graph-to-sequence pipeline. It introduces a six-step architecture—tokenization, heterogeneity alignment, inner/outer ordering, and updating via Mamba blocks—to capture both within-type and cross-type dependencies with linear-time efficiency. Empirically, HGMN outperforms 19 baselines on public heterogeneous benchmarks, achieving superior accuracy and favorable efficiency, and demonstrates robustness on both standard and long-range heterogeneous graphs. The work advances scalable, expressive heterogeneous graph learning and suggests a promising direction for integrating data-dependent state-space dynamics into graph representations.

Abstract

We propose a heterogeneous graph mamba network (HGMN) as the first exploration in leveraging the selective state space models (SSSMs) for heterogeneous graph learning. Compared with the literature, our HGMN overcomes two major challenges: (i) capturing long-range dependencies among heterogeneous nodes and (ii) adapting SSSMs to heterogeneous graph data. Our key contribution is a general graph architecture that can solve heterogeneous nodes in real-world scenarios, followed an efficient flow. Methodologically, we introduce a two-level efficient tokenization approach that first captures long-range dependencies within identical node types, and subsequently across all node types. Empirically, we conduct comparisons between our framework and 19 state-of-the-art methods on the heterogeneous benchmarks. The extensive comparisons demonstrate that our framework outperforms other methods in both the accuracy and efficiency dimensions.

HeteGraph-Mamba: Heterogeneous Graph Learning via Selective State Space Model

TL;DR

HeteGraph-Mamba (HGMN) addresses the challenge of learning from large heterogeneous graphs by marrying selective state space models with a graph-to-sequence pipeline. It introduces a six-step architecture—tokenization, heterogeneity alignment, inner/outer ordering, and updating via Mamba blocks—to capture both within-type and cross-type dependencies with linear-time efficiency. Empirically, HGMN outperforms 19 baselines on public heterogeneous benchmarks, achieving superior accuracy and favorable efficiency, and demonstrates robustness on both standard and long-range heterogeneous graphs. The work advances scalable, expressive heterogeneous graph learning and suggests a promising direction for integrating data-dependent state-space dynamics into graph representations.

Abstract

We propose a heterogeneous graph mamba network (HGMN) as the first exploration in leveraging the selective state space models (SSSMs) for heterogeneous graph learning. Compared with the literature, our HGMN overcomes two major challenges: (i) capturing long-range dependencies among heterogeneous nodes and (ii) adapting SSSMs to heterogeneous graph data. Our key contribution is a general graph architecture that can solve heterogeneous nodes in real-world scenarios, followed an efficient flow. Methodologically, we introduce a two-level efficient tokenization approach that first captures long-range dependencies within identical node types, and subsequently across all node types. Empirically, we conduct comparisons between our framework and 19 state-of-the-art methods on the heterogeneous benchmarks. The extensive comparisons demonstrate that our framework outperforms other methods in both the accuracy and efficiency dimensions.
Paper Structure (24 sections, 10 equations, 2 figures, 2 tables)

This paper contains 24 sections, 10 equations, 2 figures, 2 tables.

Figures (2)

  • Figure 1: Overview. (1) Input graph is sequentialized into tokens, each being a graph of metapath instances centered on a target node $\mathcal{V}_i$. (2a) Nodes of node type $A \in \mathcal{A}$ are projected onto the same latent representation space with a type-specific linear transformation. (2b) Node representations within each metapath instance $m$, which is a node sequence $M(s,e)$ from start node $\mathcal{V}_s$ to end node $\mathcal{V}_e$, are aggregated including the intermediate nodes. An instance encoder is used to produce vector representations $h_{M(s,e)}$ for each $m$, embedding in-between context as well. For every $\mathcal{V}_s$, the significance of each $m \in M$ is modeled through a graph attention layer to form vector representation $h_s^{M_i}$. (2c) All $h_s^{M_i}$ are aggregated to capture varying contributions of each $M \in \mathcal{M}$. (3) Nodes grouped by node type are ordered in increasing numbers of $m$. (4) Context-aware filtering is applied to the groups of nodes with a MAMBA layer for each. (5) Nodes are re-ordered in increasing order of m across all $A$. (6) Finally, a single MAMBA layer is applied to all nodes.
  • Figure 2: Time and Memory The area of the circles represents the (relative) memory consumption.

Theorems & Definitions (4)

  • Definition 2.1
  • Definition 2.2
  • Definition 2.3
  • Definition 2.4