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

## Abstract

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.

## 1 Introduction

This work proposes a heterogeneous graph mamba network (HGMN) to explore the next-generation heterogeneous graph learning by leveraging the powerful selective state space models (SSSMs), tailored to surpass the most popular transformer-based methods. This work addresses two major challenges in devising SSSMs-enhanced solution that needs greater expressiveness and minimized inference time for heterogeneous graph tasks: (i) long-range dependencies, for example, IMDB’s sparse network of 21K nodes with only 87K edges requires leveraging distant neighbor information to enhance node embeddings— a challenge amplified by the heterogeneity of the graph and (ii) graph-to-sequence conversion, where the process maps unordered graph data into a sequential structure, leveraging heterogeneous graph characteristics for effective SSSM processing.

To capture the long-range dependencies among graph, several studies [, ] integrate the transformer architecture, which utilizes full global attention to enhance the representation of diverse node and edge types. HINormer[] further extends this approach by proposing a global-range attention mechanism coupled with a dual-encoder to manage heterogeneous and structural data within graph representations. However, models employing global attention mechanisms always require a quadratic time complexity of $O(n^{2})$, which limits their scalability. To address it, models like Exphormer [] employ sparse attention techniques, such as random subsampling, to reduce computational demands, though they still face significant challenges in robustness. Recently, Mamba [], an enhanced state space model(SSM), introduces a data-dependent state transition mechanism that not only captures long-range context but also demonstrates linear-time efficiency and competitive performance with traditional Transformers. This innovation inspires researchers to adapt its architecture on many domains [, , ] and get surprising performance. Additionally, some studies [, ] propose effective strategies for mapping graph structures into ordered sequences for integrating Mamba block. Despite these achievements, the complexity of real-world heterogeneous graph scenarios remains a challenge, prompting us to propose a novel graph-to-sequence conversion mechanism to better capture long-range heterogeneity dependencies among such graphs.

Figure: 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.
Refer to caption: extracted/2405.13915v1/framework3.png

To address the above challenges, we introduce HGMN, a heterogeneous graph mamba network that provides a data-dependent alternative for attention sparsification, capable of capturing long-range heterogeneity dependencies and reducing computational costs in large heterogeneous graphs. As depicted in Figure 1, HGMN features a six-step scalable, flexible, and powerful architecture:

- 1.
Tokenization: Map a graph into a sequence of tokens, adapting sequential encoders for graphs, where each token is a subgraph of a target node and its meta-path instances.
- 2.
Heterogeneity Alignment: Project different node types into a unified latent representation space, and each token (subgraph) is aggregated into an updated target node representation.
- 3.
Inner Ordering: Group nodes by node type and order within each group by the number of meta-path instances, reflecting the importance of the nodes in the same node type.
- 4.
Inner Updating: Scans and selects relevant nodes for updating in each group of nodes to get the long-range homogeneity dependencies by a Mamba layer.
- 5.
Outer Ordering: Order all types of nodes together by node degree, reflecting the global importance of nodes in a graph.
- 6.
Outer Updating: Apply the updating process across all node types to ensure comprehensive updates and capture the long-range heterogeneity dependencies by a Mamba layer.

In the end, HGMN produces updated node embeddings for downstream tasks such as node prediction.

In summary, HGMN is the first model to capture long-range heterogeneity dependencies among large heterogeneous graphs with the support of enhanced data-dependent SSMs, delivering linear-time efficiency and satisfying performance. Our main contributions are as follows:

- •
We introduce HGMN, a novel architecture that integrates the powerful capabilities of SSSMs for addressing the unique challenges of heterogeneous graph data, particularly in improving expressiveness and minimizing inference time.
- •
We propose a unique graph-to-sequence conversion mechanism that transforms unordered graph data into a sequential format, leveraging the inherent characteristics of heterogeneous graphs to facilitate effective SSSM processing.
- •
We design a six-step scalable recipe for HGMN, which includes Tokenization, Heterogeneity Alignment, Inner and Outer Ordering, and Inner and Outer Updating steps to manage and update node representations, capturing both homogeneity and heterogeneity dependencies.
- •
Experimental results demonstrate that our framework can outperform existing transformer-based and sparse attention methods on public benchmarks, achieving higher accuracy and efficiency in public benchmarks including normal and large heterogeneous graphs.

## 2 Related Work

### 2.1 Homogeneous Graph Neural Networks

Homogeneous graph neural networks propose a massage-passing mechanism that aims to process the unstructured data (graph). GCN[] applies convolution directly on graphs to efficiently capture node interrelations for node classification and link prediction. GAT[] adds attention mechanisms to dynamically prioritize and aggregate neighboring node features in graph networks. To deal with large-scale graphs, GraphSAGE[] first uses neighborhood sampling and aggregation. While these methods can get satisfying performance for some public benchmarks, the real world has more complex heterogeneous graph structures than homogeneous ones.

### 2.2 Heterogeneous Graph Neural Networks

To solve heterogeneous edges, RGCN[] adapts GCNs for multi-relational graphs by assigning distinct weights to different relationship types. RSHN[] embeds both nodes and edges by integrating a Coarsened Line Graph with a HGNN[], capturing deep relational structures in large-scale networks. HetSANN[] directly encodes HGNNs’ structural information using an attention mechanism, bypassing the need for metapath based preprocessing.
NARS[] enhances feature smoothing on heterogeneous graphs (HGs) by averaging features over relation-specific subgraphs for scalable and memory-efficient graph learning. MAGNN[] refines HG embedding by aggregating complex relational paths and node content.
HetGNN[] and HAN[] apply hierarchical attention to enhance node embedding by utilizing metapath based neighbors. DiffMG[] and PMMM[] optimize the HGNN architectures by searching for flexible and stable meta-multigraphs that capture complex semantic relations. SeHGNN[] simplifies the HGNN architectures through a single-layer, metapath extended structure.
However, addressing long-range dependencies in sparse HG like IMDB (21K nodes with only 87K edges), which demands leveraging distant neighbor information efficiently amidst heterogeneity, remains a significant challenge.

### 2.3 Graph Transformer

To solve the limitations of Message-Passing-based model such as expressiveness bounds, over-smoothing, over-squashing and so on, some recent work employ full global attention to enhance performance, motivated by Transformer[].
GTN[] generates and utilizes new graph structures for dynamic node representation learning in heterogeneous environments without pre-defined metapaths. HGT[] leverages a transformer architecture to manage and enhance the representation of diverse node and edge types in heterogeneous graphs.
HINormer[] utilizes a global-range attention mechanism and dual-encoder modules to effectively manage heterogeneous and structural data in network representations. However, to achieve high expressive power, models employing global attention mechanisms like the Transformer require a time complexity of $O(n^{2})$. Even though sparse attention techniques can reduce this complexity, methods like Exphormer [] that use random node subsampling still face significant robustness challenges. Therefore, introducing methods capable of naturally informed context selection might present a more viable solution.

### 2.4 Selective State Space Model

State space models (SSMs) [, ] recently emerge as a popular alternative to attention-based modeling architectures due to their efficiency, which stems from performing recurrent updates across sequences via hidden states. However, their effectiveness is often limited compared to Transformers because of their time-invariant transition mechanisms. To address this, Mamba [] introduces a data-dependent state transition mechanism that captures long-range context, demonstrating linear-time efficiency and competitive performance with Transformers. This innovation motivates researchers to adapts its architecture across many domains [, , ]. In graph domain, some studies [, ] propose effective strategies for mapping graph structures into ordered sequences. Others [, ] develop frameworks for spacial-temporal graph and point cloud tasks, showcasing the efficacy of Mamba-based models in specific benchmarks. However, the real world presents a more complex array of heterogeneous graph scenarios. It motivates us to propose a graph-to-sequence conversion mechanism to capture long-range heterogeneity dependencies among heterogeneous graphs.

### 2.5 Preliminaries

State Space Models

The SSM-based models are inspired by the continuous system, which maps a sequence $x(t)\in\mathbb{R}\mapsto y(t)\in\mathbb{R}$ through a hidden state $h(t)\in\mathbb{R}^{\mathtt{N}}$. This system uses $\mathbf{A}\in\mathbb{R}^{\mathtt{N}\times\mathtt{N}}$ as the updating parameter and $\mathbf{B}\in\mathbb{R}^{\mathtt{N}\times 1}$, $\mathbf{C}\in\mathbb{R}^{1\times\mathtt{N}}$ as the projection parameters.

$$ $\displaystyle h^{\prime}(t)=\mathbf{A}h(t)+\mathbf{B}x(t),y(t)=\mathbf{C}h(t).$ (1) $$

For example, S4 [] and Mamba [] models are discrete versions of the system, incorporating a timescale parameter $\mathbf{\Delta}$ that converts the continuous coefficients $\mathbf{A}$ and $\mathbf{B}$ into discrete equivalents $\mathbf{\overline{A}}$ and $\mathbf{\overline{B}}$. The transformation employs the zero-order hold (ZOH) technique, defined by the equations:

$$ $\displaystyle\mathbf{\overline{A}}=\exp{(\mathbf{\Delta}\mathbf{A})},\mathbf{ \overline{B}}=(\mathbf{\Delta}\mathbf{A})^{-1}(\exp{(\mathbf{\Delta}\mathbf{A} )}-\mathbf{I})\cdot\mathbf{\Delta}\mathbf{B}.$ (2) $$

Following this discretization process, the discrete form of Eq. (1) with a step interval of $\mathbf{\Delta}$ is as:

$$ $\displaystyle h_{t}=\mathbf{\overline{A}}h_{t-1}+\mathbf{\overline{B}}x_{t},y_ {t}=\mathbf{C}h_{t}.$ (3) $$

Finally, the models produce the output by executing a global convolution:

$$ $\displaystyle\mathbf{\overline{K}}=(\mathbf{C}\mathbf{\overline{B}},\mathbf{C} \mathbf{\overline{A}}\mathbf{\overline{B}},\dots,\mathbf{C}\mathbf{\overline{A }}^{\mathtt{M}-1}\mathbf{\overline{B}}),\mathbf{y}=\mathbf{x}*\mathbf{ \overline{K}},$ (4) $$

where $\mathtt{M}$ denotes the sequence length $\mathbf{x}$, and $\overline{\mathbf{K}}\in\mathbb{R}^{\mathtt{M}}$ represents a structured convolutional kernel.

###### Definition 2.1 .

Heterogeneous Graph.
We define a heterogeneous graph as $\mathcal{G}=(\mathcal{V},\mathcal{E})$, where $\mathcal{V}$ and $\mathcal{E}$ denote the sets of nodes and edges, respectively. Each node and edge is associated with predefined types through mapping functions: $\mathcal{V}\rightarrow\mathcal{A}$ for nodes and $\mathcal{E}\rightarrow\mathcal{R}$ for edges.

###### Definition 2.2 .

Metapath.
We define a metapath $M$ of its set $\mathcal{M}$ as $A_{1}\stackrel{{\scriptstyle R_{1}}}{{\longrightarrow}}A_{2}\stackrel{{
\scriptstyle R_{2}}}{{\longrightarrow}}\cdots\stackrel{{\scriptstyle R_{i}}}{{
\longrightarrow}}A_{i+1}$ including a composite relation $R=R_{1}\rhd R_{2}\rhd\cdots\rhd R_{i}$ that covers the relationships between the starting node type $A_{1}$ and the ending node type $A_{i+1}$, where the $\rhd$ symbol denotes the composition.

###### Definition 2.3 .

Metapath Instance.
We define a metapath instance $m$ of given metapath $M$ as a node sequence in the graph following the schema defined by $M$.

###### Definition 2.4 .

Metapath Graph.
We define a metapath graph $\mathcal{G}^{M}$ as a subgraph constructed by nodes from all metapath-M-based node sequences in graph $\mathcal{G}$.

###### Definition 2.1 .

###### Definition 2.2 .

###### Definition 2.3 .

###### Definition 2.4 .

## 3 Methodology

### 3.1 Overview

In the Tokenization step, we map the graph into a sequence of tokens, adapting sequential encoders for graph data. Each token represents a subgraph that includes a target node and its all metapath instances, where a metapath instance is an ordered sequence of nodes within a composite relation among the node types involved.
Subsequently, in the Heterogeneity Alignment step, we use a mapping function to project nodes of different types into the same latent representation space, aggregating each token (subgraph) into an updated target node representation.
Following this, in the Inner Ordering step, nodes are grouped by node type and ordered within each group based on the number of the metapath instances. This ordering reflects an assumption that the greater the number of metapath instances associated with a node, the more significant that node is within its type.
In the Inner Update step, a Mamba mechanism is employed to scan and select relevant nodes for updating. Due to the recurrent updates, each token accesses information only from preceding tokens. Consequently, more important nodes are positioned closer to the end of the sequence to maximize their visibility and impact.
Upon completing the updates within each node type, in the Outer Update step, all nodes are ordered by their degree, and a similar updating process is applied across all node types to ensure global updates.
Ultimately, HGMN outputs the final node embeddings, which are ready for downstream tasks such as node prediction, ensuring that the system captures the long-range heterogeneity dependencies.

### 3.2 Tokenization

To adapt the sequential SSM-based model for graph data, we need a process to divide a given graph into many tokens. Recent works try to employ node, edge, or even subgraph tokenization methods, each with its own advantages and disadvantages [, , ]. In this step, we design a simple but effective and efficient strategy that considers both the heterogeneity and homogeneity of graphs.

Given a heterogeneous graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$, with nodes and edges mapped to predefined types through the functions $\mathcal{V}\rightarrow\mathcal{A}$ and $\mathcal{E}\rightarrow\mathcal{R}$, we generate a subgraph as a token for each node $\mathcal{V}_{i}$:

$$ $\mathcal{G}[\mathcal{V}_{i}]=\bigcup_{k=1}^{|\mathcal{M}|}\mathcal{G}[\mathcal {V}_{i}]^{M_{k}}$ (5) $$

Each subgraph $\mathcal{G}[\mathcal{V}_{i}]$ includes a target node $\mathcal{V}_{i}$ and its metapath instances from all metapaths in the set $\mathcal{M}$. After tokenization, instead of using individual nodes, we utilize a subgraph $\mathcal{G}[\mathcal{V}_{i}]$ as the representative token for each node.

### 3.3 Heterogeneity Alignment

To aggregate diverse information, this step aligns the heterogeneity within each token (subgraph).

Type-specific Linear Transformation:
We first address the diversity of node types by using a type-specific linear transformation to project feature vectors of different node types into the same latent representation space. This is formalized as follows for each node type $A\in\mathcal{A}$:

$$ $\mathbf{h}_{i}^{A}=\mathbf{W}_{A}\mathbf{x}_{i}^{A}+\mathbf{b}_{A},$ (6) $$

where $\mathbf{x}_{i}^{A}\in\mathbb{R}^{d_{A}}$ is the feature vector of node $\mathcal{V}_{i}$ of type $A$, and $\mathbf{W}_{A}\in\mathbb{R}^{d^{\prime}\times d_{A}}$ and $\mathbf{b}_{A}\in\mathbb{R}^{d^{\prime}}$ are the parameters of the linear transformation specific to type $A$.

Metapath Instance Aggregation: Following the projection, we then start to aggregate each token (subgraph) by two levels: (1) metapath instance level and (2) metapath level.

Within a metapath instance $m$ of metapath $M$, there is node sequence $M(s,e)$ including a start node $\mathcal{V}_{s}$, an end node $\mathcal{V}_{e}$, and the intermediate nodes. We use an instance encoder to transfer all the node features within $M(s,e)$ to a single vector: $\mathbf{h}_{M\left(s,e\right)}=f_{\theta}\left(M(s,e)\right)=f_{\theta}\left(
\left\{\mathbf{h}_{t}^{\prime},\forall t\in\{M(s,e)\}\right\}\right)$
where $\mathbf{h}_{M\left(s,e\right)}\in\mathbb{R}^{d^{\prime}}$ has a dimension of $d^{\prime}$.

After encoding metapath instances, we utilize the graph attention layer to obtain the weighted sum of all metapath instances for each metapath $M$ related to targe node $\mathcal{V}_{s}$. The primary rationale is that each instance contributes to the representation of $\mathcal{V}_{s}$ to varying degrees. To model the significance of each metapath instance through importance weights $\alpha_{s,e}^{M}$, the graph attention layer firstly computes $e_{se}^{M}$ (quantifying the relevance of the metapath instance $M(s,e)$ to node $\mathcal{V}_{s}$) by $e_{se}^{M}=\operatorname{LeakyReLU}\left(\mathbf{a}_{M}\cdot\left[\mathbf{h}_{
s}^{\prime}\|\mathbf{h}_{M\left(s,e\right)}\right]\right)$, where
the attention vector $\mathbf{a}_{M}$, specific to metapath $M$ and parameterized in $\mathbb{R}^{2d^{\prime}}$, facilitates the use of the vector concatenation operator, $\|$. Then, it models the significance of each metapath instance through importance weights $\alpha_{s,e}^{M}$:

$$ $\alpha_{se}^{M}=\frac{\exp\left(e_{se}^{M}\right)}{\sum_{e\in\mathcal{N}_{s}^{ M}}\exp\left(e_{se}^{M}\right)}.$ (7) $$

This significance is normalized across all possible $\mathcal{V}_{e}$ in $\mathcal{N}{s}^{M}$ using the softmax function. Following normalization, the importance weights $\alpha_{s,e}^{M}$ for each $e$ in $\mathcal{N}_{s}^{M}$ are used to form a weighted sum of the representations from the metapath instances associated with node $\mathcal{V}_{s}$). Finally, we aggregate these contributions by summing all instances with an activation function $\sigma(\cdot)$:

$$ $\displaystyle\mathbf{h}_{s}^{M}=\sigma\left(\sum_{e\in\mathcal{N}_{s}^{M}} \alpha_{se}^{M}\cdot\mathbf{h}_{M\left(s,e\right)}\right).$ (8) $$

Metapath Aggregation:
After aggregating the representations from all instances within each metapath, we proceed to aggregate the outputs across different metapaths for the target node $\mathcal{V}_{s}$. Assuming $\mathcal{V}_{s}$ is of type $A$, we gather a set of latent vectors: $\{\mathbf{h}_{s}^{M_{1}},\mathbf{h}_{s}^{M_{2}},\ldots,\mathbf{h}_{s}^{M_{K}}\}$, where $K$ denotes the number of metapaths associated with $\mathcal{V}_{s}$. A straightforward approach would be using the average vector. However, to capture the distinct contributions of different metapaths, we employ a similar mechanism to the previous aggregation, enhancing the representational power of the aggregation:

$$ $\mathbf{h}_{s}^{\prime}=\sum_{k=1}^{K}\beta_{s}^{M_{k}}\cdot\mathbf{h}_{s}^{M_ {k}},$ (9) $$

where importance weights $\beta_{s}^{M_{k}}$ are used to compute a weighted sum of the metapath representations, thereby finalizing the aggregated output for node $\mathcal{V}_{s}$. This aggregation synthesizes information across different metapaths, allowing the model to highlight more significant metapaths based on their relevance to the target node’s type and context. The resulting vector $\mathbf{h}_{s}^{\prime}$ represents a comprehensive embedding of $\mathcal{V}_{s}$, considering both its type-specific properties and connectivity in the graph.

After two-level aggregation, each token becomes an updated representation $\left\{\mathbf{h}_{i}^{\prime},\forall i\in\{\mathcal{V}\}\right\}$.

### 3.4 Inner Ordering & Updating

After aggregating each token into an updated representation, we start to capture long-range heterogeneity and homogeneity dependencies among the graph.
Firstly, all nodes are grouped by their types and ordered within each group based on the number of metapath instances. This ordering reflects an assumption that the greater the number of metapath instances associated with a node, the more significant that node is within its type. After that, each Mamba block updates node representations within each node type following the ordered sequence.

#### 3.4.1 Inner-Type Ordering

The Need for Ordering in Graphs:
Graphs are unordered structures, where the arrangement of nodes does not imply any specific processing sequence. However, for SSM-based models such as Mamba, which are adapted from architectures designed for sequential data, converting the graph data into an ordered sequence is imperative. Ordering nodes provide a structured way to input data into these models, enabling the application of advanced architectures developed for ordered data.

Efficient Ordering Method:
Given the unordered graph data, proposing an efficient method to order nodes is essential. This strategy addresses it by grouping nodes according to type and then ranking them within their groups based on their number of metapath instances:

$$ $\mathcal{V}_{A}=\{\mathcal{V}_{i}\in\mathcal{V}:\text{type}(\mathcal{V}_{i})=A \},M_{i}=\sum_{M\in\mathcal{M}}\text{count}(\mathcal{V}_{i},M).$ (10) $$

Nodes $\mathcal{V}_{i}$ in each group $\mathcal{V}_{A}$ are ordered in ascending order of $M_{i}$. Because each node in the Mamba block can only gather information from previous nodes, the most influential nodes are placed towards the end of the sequence, ensuring that they have the maximal possible context for their feature representations, enhancing the model’s effectiveness. This metric not only reflects the quantitative involvement of nodes in graph structures but also prioritizes nodes with higher connectivity and semantic importance within the graph. This prioritization is aligned with the model’s need to capture the most relevant information early in its operations.

By ordering the nodes in this manner, we harness the complex and rich information present in heterogeneous graphs and adapt it for models requiring sequential data input, thereby bridging the gap between unordered graph data and the ordered input requirements of SSMs.

#### 3.4.2 Inner-Type Updating

Context-Aware Filtering in Mamba:
Mamba optimizes the processing of graph data by integrating a selective filtering mechanism that distinguishes between relevant and irrelevant contexts. This functionality is critical for managing the extensive and varied connections typical in graphs, enabling the model to maintain focus on pertinent information over extended sequences.

Mechanism for Dynamic Information Filtering:
Mamba employs the matrices $\mathbf{\overline{A}}$ and $\mathbf{\overline{B}}$ to modulate the influence of past states and current inputs:
$\mathbf{h}_{t}=\mathbf{\overline{A}}h_{t-1}+\mathbf{\overline{B}}x_{t}.$

These matrices enable the model to ’remember’ or ’forget’ information, facilitating the management of long-range dependencies where the relevance of information varies across the graph.

Sparsification of Graph Attention:
The model enhances traditional graph attention mechanisms by sparsifying attention, focusing computational resources on the graph’s most informative parts to reduce complexity and enhance performance:
$\mathbf{y}_{t}=\mathbf{C}h_{t},\mathbf{y}=\mathbf{x}*\mathbf{\overline{K}}$,
where $\mathbf{\overline{K}}$ aggregates significant features across the graph, implementing context-aware filtering.

Implications: (1) Enhances efficiency: Reduces computational overhead by focusing on critical interactions; (2) Improves accuracy: Decreases noise in data, increasing prediction accuracy; (3) Preserves long-range dependencies: Ensures critical nodes are maintained over large graph.

Overall, Mamba’s updating for context-aware filtering enables better management of graph data by focusing on relevant features, which is crucial for handling complex network structures.

### 3.5 Outer Ordering & Updating

After utilizing $|\mathcal{A}|$ Mamba to update nodes from $|\mathcal{A}|$ distinct node types in the previous step, we capture the long-range homogeneity dependencies within each node type. Subsequently, our focus shifts to addressing the long-range heterogeneity dependencies that exist across different node types. We design a new Ordering & Updating strategy for this purpose, as detailed below.

#### 3.5.1 Cross-Type Ordering

Extending Ordering Across Node Types:
Having established an order within individual node types, we now extend it to span across different node types. This outer ordering is crucial for capturing heterogeneous long-range dependencies reflecting global structural patterns within the graph.

Cross-Type Ordering Strategy Based on Degree:
To achieve this ordering, we prioritize nodes based on node degree. An overarching sequence is then determined, which indicates their global influence and connectivity. This sequence ensures that nodes acting as critical hubs between different node types are positioned to maximize their contextual influence throughout the graph.

Implementation of Heterogeneous Dependencies:
By arranging nodes to maximize within-type relationships and then leverage between-type interactions, we capture the full spectrum of dependencies within the graph. It enhances the model’s ability to process heterogeneous data and aligns with Mamba’s capability to handle complex inter-dependencies inherent in large-scale graph structures.

#### 3.5.2 Cross-Type Updating

Global Context-Aware Updating:
With the outer ordering, updating node representations across different types involves an approach to integrate and filter information from global sources. This step is essential for managing the diversity of interactions among heterogeneous node types.

Dynamic Cross-Type Information Filtering:
Utilizing structured matrices $\mathbf{\overline{A}}$ and $\mathbf{\overline{B}}$, the model now updates each node’s state considering its degree and type-specific context:
$\mathbf{h}_{t}^{global}=\mathbf{\overline{A}}h_{t-1}^{global}+\mathbf{
\overline{B}}x_{t}^{cross},$ where $x_{t}^{cross}$ represents the aggregated input features from multiple node types, emphasizing the integration of diverse information streams.

Enhanced Graph Attention for Heterogeneity:
Attention mechanisms are adapted to dynamically allocate focus across different node types, balancing intra-type precision with inter-type coverage: $\mathbf{y}_{t}=\mathbf{C}h_{t}^{global},\mathbf{y}=\mathbf{x}*\mathbf{
\overline{K}}^{global},$ where $\mathbf{\overline{K}}^{global}$ extends the convolutional kernel to encapsulate interactions across node types, refining the model’s ability to synthesize and highlight key features from a graph-wide perspective.

Outcomes: Updated nodes’ representations by combining information from all node types.

## 4 Experiments

In this section, we compare our model with many state-of-the-art baselines on standard heterogeneous graph benchmark (HGB) [] and long-range heterogeneous graph benchmark (LRHGB) [].

### 4.1 Datasets and Baselines

Long Range Heterogeneous Graph Benchmark

We select the large-scale dataset ogbn-mag from OGB challenge []. ogbn-mag is a heterogeneous graph derived from a subset of the Microsoft Academic Graph (MAG). It comprises four kinds of entities — papers, authors, institutions, and fields of study — along with four types of directed relationships linking two different entity types.

Standard Heterogeneous Graph Benchmark

We choose three heterogeneous graphs including DBLP, IMDB, and ACM from HGB benchmark [] on the node classification task. DBLP is an online bibliography resource for computer science. We employ a frequently utilized subset spanning four domains, where nodes represent authors, papers, terms, and venues. IMDB is a portal dedicated to movies and associated details. We utilize a subset covering genres such as Action, Comedy, and Drama. ACM is a citation network as well. We utilize the subset provided in HAN [], maintaining all connections including paper citations and references.

We choose MLP[], GraphSAGE[], RGCN[], NARS[], HGT[], SeHGNN[], LMSPS[], RSHN[], HetSANN[], GTN[], HGT[], Simple-HGN[], HINormer[], RGCN[], HetGNN[], HAN[], MAGNN[], SeHGNN[], DiffMG[], PMMM[], and Graph-Mamba [] (adding node-type linear projection at the head) as our baselines.

**Table 1: Performance of Node Classification We conduct an experiment using 19 baselines of heterogeneous graph modeling methods, evaluated with the F1 score and accuracy. Across all configurations, HeteGraph-Mamba achieves the best performance in the obgn-mag, DBLP, IMDB and ACM datasets. Highlighted are the top first, second, third.**
| Model | ogbn-mag | DBLP | IMDB | ACM |
| --- | --- | --- | --- | --- |
|  | Accuracy | F1 score | F1 score | F1 score |
| MLP | 0.2692$\pm$0.0026 | 0.4933$\pm$0.0030 | 0.2815$\pm$0.0033 | 0.4018$\pm$0.0104 |
| GCN | 0.3489$\pm$0.0019 | 0.9084$\pm$0.0032 | 0.5273$\pm$0.0024 | 0.8743$\pm$0.0135 |
| GAT | 0.3767$\pm$0.0025 | 0.9142$\pm$0.0027 | 0.5380$\pm$0.0094 | 0.8845$\pm$0.0136 |
| GraphSAGE | 0.4678$\pm$0.0067 | 0.9182$\pm$0.0011 | 0.5591$\pm$0.0036 | 0.8983$\pm$0.0098 |
| RGCN | 0.4737$\pm$0.0048 | 0.9207$\pm$0.0050 | 0.6205$\pm$0.0015 | 0.9141$\pm$0.0075 |
| RSHN | 0.4728$\pm$0.0031 | 0.9381$\pm$0.0055 | 0.6422$\pm$0.0103 | 0.9032$\pm$0.0154 |
| HetSANN | 0.4781$\pm$0.0029 | 0.8056$\pm$0.0150 | 0.5768$\pm$0.0044 | 0.8991$\pm$0.0037 |
| NARS | 0.5088$\pm$0.0012 | 0.8845$\pm$0.0034 | 0.6539$\pm$0.0018 | 0.9310$\pm$0.0040 |
| MAGNN | 0.4926$\pm$0.0024 | 0.9376$\pm$0.0045 | 0.6467$\pm$0.0167 | 0.9077$\pm$0.0065 |
| HetGNN | 0.4732$\pm$0.0130 | 0.9176$\pm$0.0043 | 0.4825$\pm$0.0067 | 0.8591$\pm$0.0025 |
| HAN | 0.5037$\pm$0.0066 | 0.9205$\pm$0.0062 | 0.6463$\pm$0.0058 | 0.9079$\pm$0.0043 |
| DiffMG | 0.5157$\pm$0.0044 | 0.9420$\pm$0.0036 | 0.5975$\pm$0.0123 | 0.8807$\pm$0.0304 |
| PMMM | 0.5188$\pm$0.0031 | 0.9514$\pm$0.0022 | 0.6758$\pm$0.0022 | 0.9371$\pm$0.0017 |
| Simple-HGN | 0.5215$\pm$0.0013 | 0.9446$\pm$0.0022 | 0.6736$\pm$0.0057 | 0.9335$\pm$0.0045 |
| SeHGNN | 0.5671$\pm$0.0014 | 0.9524$\pm$0.0013 | 0.6821$\pm$0.0032 | 0.9367$\pm$0.0050 |
| GTN | 0.5392$\pm$0.0037 | 0.9397$\pm$0.0054 | 0.6514$\pm$0.0045 | 0.9120$\pm$0.0071 |
| HGT | 0.4929$\pm$0.0061 | 0.9349$\pm$0.0025 | 0.6720$\pm$0.0057 | 0.9100$\pm$0.0076 |
| HINormer | 0.5520$\pm$0.0062 | 0.9494$\pm$0.0021 | 0.6783$\pm$0.0034 | 0.9379$\pm$0.0122 |
| Graph-Mamba++ | 0.5026$\pm$0.0020 | 0.9271$\pm$0.0019 | 0.6322$\pm$0.0096 | 0.9017$\pm$0.0108 |
| HGMN | 0.5763$\pm$0.0043 | 0.9602$\pm$0.0010 | 0.6917$\pm$0.0028 | 0.9484$\pm$0.0046 |
| - w/o Inner-ordering | 0.5371$\pm$0.0034 | 0.9410$\pm$0.0041 | 0.6535$\pm$0.0079 | 0.9112$\pm$0.0130 |
| - w/o Outer-ordering | 0.5237$\pm$0.0040 | 0.9342$\pm$0.0031 | 0.6492$\pm$0.0036 | 0.9209$\pm$0.0085 |

### 4.2 Benchmark Evaluation

Comparison on Performance

We evaluate the performance of our HGMN against 19 heterogeneous graph baselines using accuracy and F1 score metrics across several datasets shown in Table 1. Our model demonstrated superior performance, achieving the highest scores on the ogbn-mag, DBLP, and ACM datasets with an accuracy of $0.5763\pm 0.0043$ on ogbn-mag, and F1 scores of $0.9602\pm 0.0010$ on DBLP, and $0.9484\pm 0.0046$ on ACM. It also performed strongly on IMDB with an F1 score of $0.6917\pm 0.0028$. These results highlight HGMN’s robustness and its ability to manage complex graph structures, validating its advanced node classification capabilities in heterogeneous graphs.

Comparison on Efficiency

In our study, we evaluate the time and memory requirements of various HGNNs on the DBLP dataset shown in Figure 2. This analysis highlights our model’s efficiency, achieving an excellent F1 score with moderate memory use and computational time. Other models like SeHGNN and HINormer, although high-performing, required significantly more memory. This evaluation emphasizes our model’s capability to efficiently process large-scale graphs without compromising performance.

Figure: Figure 2: Time and Memory The area of the circles represents the (relative) memory consumption.
Refer to caption: extracted/2405.13915v1/efficiency.png

## 5 Summary

In this study, we present HGMN, a cutting-edge network that utilizes SSSMs to redefine heterogeneous graph learning, outperforming transformer-based approaches. Our strategy transfers unordered graph data into a two-level sequential format. The first level organizes graph data to capture homogeneity dependencies among similar node types and subsequently leverages organized sequences to address heterogeneity dependencies across different types of nodes. This dual sequential process not only enhances the expressiveness and computational efficiency of the linear-time performance, but also improves the accuracy of node embeddings for downstream tasks. Our experimental results validate HGMN’s superior capability in managing the complexities of large heterogeneous graphs, demonstrating marked performance improvements over existing methods on benchmarks. HGMN also has limitations on theoretical analysis, which can provide more strong proofs.

## Broader Impact

Our research methodology can bolster comprehension and problem resolution across numerous areas, including AI research, fostering clearer and more decipherable outcomes. However, this method might simplify intricate problems too much by dividing them into distinct segments, possibly neglecting subtleties and linked components. Moreover, a strong dependence on this approach could curtail innovative problem-solving, since it promotes a sequential and orderly method, which may hinder unconventional thought processes.

## Appendix A Supplemental Material

### A.1 System

All experiments are carried out on a High Performance Computing cluster. There are 34 GPU nodes where 16 nodes each have 2 NVIDIA 40GB Tesla A100 PCIe GPUs, 52 CPU cores, and 192 GB of CPU RAM while 18 nodes are each equipped with 4 NVIDIA 80GB Tesla A100 SXM GPUs, 52 CPU cores, and 512 GB of CPU RAM. The driver version 525.105.17 on these nodes is compatible with CUDA 12.0 or earlier. The operating system is Red Hat Enterprise Linux 7.9.

### A.2 Hyperparameter

In our experiments, we require hyperparameters for tokenization, heterogeneity alignment (Num heads and Metapath Attention Dim are exclusive to this step), and the MAMBA layers. We use the Adam optimizer with weight decay. Below, we provide a list of all the hyperparameters used in our experiments.

**Table 2: Hyperparameter used in the task.**
| Hyperparameter | ogbn-mag | DBLP | IMDB | ACM |
| --- | --- | --- | --- | --- |
| Num Layers | 4 | 2 | 2 | 2 |
| Hidden Dim | 128 | 64 | 64 | 64 |
| Learning Rate | 0.003 | 0.0005 | 0.0005 | 0.0005 |
| Weight Decay | 0.0005 | 0.0001 | 0.0001 | 0.0001 |
| Num Epochs | 300 | 150 | 150 | 150 |
| Num heads | 8 | 8 | 8 | 8 |
| Metapath Attention Dim | 256 | 128 | 128 | 128 |

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. How do you get the subgraph for all nodes? Does it include every node next to it?
2. Also please address the fourth point in weakness.
soundness: 2
strengths: 1. it is good to adopt MAMBA for graph data.
2. Comprehensive baseline models are used.
3. The proposed model is relatively fast empirically.
summary: This paper proposes a graph mamba network as a graph model to learn node representations, aiming to incorporate long-range dependencies for heterogeneous graphs. It adopts subgraphs containing multiple meta-path instances to be tokens of all nodes. Aggregate embeddings of meta-path instances for node $v$ within one meta-path and across multiple meta paths using attention mechanism to get the node embedding and then uses the updated method in MAMBA to update node embeddings within one node type first and then across node types to model the long-range dependency. The experimental results show marginal improvement.
weaknesses: 1. The related work lacks targeting the problem that the paper wants to solve. i.e., if you claim the proposed model can do long-range dependency, so your related work should discuss the work that is related to this challenge. 
2. The paper has typos and doesn't provide explanations for some symbols, i.e., $\mathcal{N}_s^M$, and `spacial` at line 172.
3. Additionally, the experiments are unable to be reproducible. it is unclear how metapaths are built. Since the improvement is marginal, it would be better to show the variance of the performance. This paper lacks an ablation study for the components in the model. sin
4. My major concern is that the local information of a node might gradually disappear during the inner and outer updates since the order of these two processes is only determined by the degree.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. How can you prove that your tokenization is both effective and efficient？
2. For Equation 6, is it reliable to map vectors of different types of nodes to the same latent
space using just a single linear layer? If this approach is successful, does it suggest that
the node feature vectors xi constructed in this paper already exist in a larger latent space,
and that a simple linear layer is used for dimension reduction? Could this operation
interfere with subsequently capturing homogenous and heterogeneous dependencies
within the scope?
3. In the ordering process, if multiple nodes have the same degree or the same number of
relevant metapaths, how should they be handled?
soundness: 2
strengths: 1. The proposed recipe to implement the HGMN is highy scalable and take both the long-range homogeneity and heterogeneity into take into consideration.
summary: This paper applies Mamba to heterogeneous graph learning tasks. The authors provide a full recipe to build the model, which includes tokenization, heterogeneity alignment, inner and outer ordering, and inner and outer updating steps. The six-step recipe help convert the heterogeneous graph to sequence, capture long-range dependency among various types of nodes as well as reduce computation burden. The empirical results demonstrate that the proposed HGMN enjoys higher accuracy and more efficiency than exsiting transformer-based GNNs on multiple heterogeneous graph benchmarks
weaknesses: 1. The motivation is too weak. Since Mamba claims to handle long-range context better
and operates in linear time complexity, the paper applies Mamba to heterogeneous graph
learning tasks, expecting superior results. This point lacks persuasiveness.
2. Novelty is very limited. The paper shares many similarities with  Graph Mamba（KDD 2024）.

2.1 For instance, in the Tokenization section, both fundamentally encode subgraphs to
represent nodes, with the difference being that Graph Mamba uses a random walk
approach to sample subgraphs, while this paper uses metapath instances.

2.2Additionally, the language style in the Tokenization section of the article is very similar
to that of Graph Mamba. 

2.3 In this paper, I think the only difference between it and Graph Mamba lies in the inner
and outer ordering methods. However, these ordering methods lack persuasiveness.
*Why does the inner ordering rely on the number of metapaths as its basis?
*Why can’t the number of node degree be used as a basis during inner ordering?
*Shouldn’t the importance of different metapaths be considered?

3. Some important detail  is missing. For example, in the Metapath Instance Aggregation
section, the implementation of the instance encoder is not explained. How is beta
implemented in Equation 9?

4. The experiments are not enough. The ablation experiments are not
analyzed and are not comprehensive, for example, there is a lack of ablation on the
TOKENIZATION and HETEROGENEITY ALIGNMENT sections. There is also a lack of
analysis on key parameters in the method.


Reference:*Graph Mamba: Behrouz, Ali, and Farnoosh Hashemi. "Graph mamba: Towards learning on
graphs with state space models." Proceedings of the 30th ACM SIGKDD Conference on
Knowledge Discovery and Data Mining. 2024.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: Q1: Why was Mamba chosen for graphs? Mamba shows improved efficiency over Transformers when modeling sequences longer than 2K tokens, based on current knowledge. Providing a theoretical analysis to support this choice in graph contexts would add significant value to the paper.
Q2: Are all six processing steps necessary? A more detailed ablation study on each step would better clarify the contributions and necessity of each component in the model.
Q3: Performance on tasks beyond node classification: Can the framework achieve similar high performance on other tasks beyond node classification?
soundness: 2
strengths: S1：The integration of SSSMs in heterogeneous graph learning is innovative, addressing scalability and long-range dependency challenges not fully tackled by traditional Transformer-based models. HGMN’s tokenization and ordering mechanisms for heterogeneous nodes offer a fresh approach to sequence modeling in graphs.
S2：The six-step processing flow in HGMN separately addresses homogeneity and heterogeneity dependencies, enhancing the model's expressive capability. Additionally, the intra-type and inter-type node ordering improves the accuracy of context capture.
S3：The paper provides a detailed overview of each component of the HGMN framework, making it easy to follow the flow and understand the contributions of each design choice.
S4：HGMN’s performance improvements over 19 baseline models on accuracy and efficiency metrics demonstrate the potential impact of this approach in the heterogeneous graph learning domain.
summary: The paper presents HeteGraph-Mamba (HGMN), a novel heterogeneous graph learning framework that uses Selective State Space Models (SSSMs) to address the limitations of existing graph neural networks. HGMN introduces a six-step process, including graph tokenization and context-aware node ordering, to improve computational efficiency and capture long-range dependencies across heterogeneous nodes. Experimental results indicate that HGMN achieves higher accuracy and efficiency than other state-of-the-art models on several benchmark datasets.
weaknesses: W1. The motivation for using Mamba on graphs is unclear. Mamba is faster than Transformer only when modeling sequences longer than 2k.
W2. Although the paper demonstrates the superiority of the experimental results, it lacks sufficient theoretical analysis of the proposed method, especially in the theoretical basis for improving the efficiency and accuracy of SSSMs. It is recommended to add more mathematical analysis or complexity proof.
W3. The previous experiments only verified the effect of HGMN in the node classification task. The applicability and performance of other graph tasks (such as edge prediction and graph classification) have not been tested, which affects its wide applicability.
W4. The model hyperparameters are not explained enough, and the analysis of the functions of each module is also lacking. It is recommended to increase the analysis of hyperparameters and add ablation experiments.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: Please see Weaknesses.
soundness: 2
strengths: The paper's description of the model is relatively clear. The language is fluent.
summary: This paper introduces HGMN, the first to apply selective state space models (SSSMs) to heterogeneous graph learning. HGMN captures long-range dependencies and adapts SSSMs for heterogeneous data. It outperforms 19 state-of-the-art methods in both accuracy and efficiency.
weaknesses: 1. The problem statement is unclear. The introduction does not effectively clarify whether the focus is on learning heterogeneous graph representations or capturing long-range dependencies. I suggest the authors clearly delineate the specific focus and differentiate between these two aspects in the introduction.

2. Lack of sufficient motivation. The paper introduces challenges related to selective state space models (SSSMs) but fails to explain the necessity of integrating SSSMs. I recommend the authors elaborate on the limitations of existing approaches and provide examples where SSSMs would be particularly beneficial to enhance the motivation.

3. The experimental results are not convincing:

(1) The paper mentions two recent baselines—LMSPS (Li et al., 2024a) and Graph-Mamba (Wang et al., 2024)—but LMSPS is not included in the experiments. The authors should explain why LMSPS was omitted and discuss its relevance for comparison.

(2) Graph-Mamba focuses on long-range graph sequence modeling rather than heterogeneous graph representation learning. The authors should justify the appropriateness of using it as a baseline.

(3) The paper claims that HGMN performs well on large heterogeneous graphs, yet the four datasets used are all standard datasets, not large-scale ones. I suggest the authors consider using specific large-scale heterogeneous graph datasets to better support their claims regarding HGMN's performance.