ICLR2025|北航、港理工等提出全新NodeIDs框架，學習超短離散節點表示

▲ Node ID 示意圖

我們簡要介紹了一種名為 Node Identifiers（Node IDs）的新型節點離散表示學習框架。該框架利用向量量化（Vector Quantization，VQ）技術，將經典圖神經網路（GNNs）在多層鄰域資訊聚合後生成的連續嵌入（embeddings），壓縮成高效、緊湊（通常 6-15 維）、並且可解釋的離散表示（int4 型別）。

透過對 34 個基準任務（涵蓋節點分類、圖分類、連結預測及圖聚類任務）的實證研究，Node IDs 在保證效能表現的同時顯著降低了記憶體佔用並加速了推理過程。

更重要的是，這些 Node IDs 不依賴可以無縫結合現有的無監督和有監督 GNNs 方法，從而兼顧預測效能與高效推理。本研究為圖表示學習在大規模場景下的應用提供了新的思路。

論文題目：

Node Identifiers: Compact, Discrete Representations for Efficient Graph Learning

論文作者：

羅元凱，李宏康，劉奇煚，時磊，吳曉明

作者單位：

北京航空航天大學、香港理工大學、倫斯勒理工學院

論文連結：

https://openreview.net/forum?id=t9lS1lX9FQ

程式碼連結：

https://github.com/LUOyk1999/NodeID

引言

在圖資料上進行機器學習時，通常要考慮圖的結構以及節點/邊特徵，這廣泛應用於節點分類、圖分類、連結預測及推薦系統等。

現有的大部分圖神經網路（GNNs）透過訊息傳遞（message passing）迭代聚合節點鄰域的資訊，在眾多工上取得了優異的效能。然而，GNNs 在大規模場景中的推理效率往往受限於以下方面：

訊息傳遞的高開銷：推理時需要載入完整的圖結構，對大規模圖（可能包含數十億條邊）會帶來巨大的計算與儲存壓力。
連續嵌入維度高且可解釋性差：典型的 GNN 嵌入通常維度較高（128 或 256 維），不僅影響儲存和計算效率，也在一定程度上缺乏可解釋性。

為此，我們提出了一個端到端的向量量化框架 NID，可在自監督或有監督訓練過程中，不引入重構損失的前提下，將多層 GNN 嵌入直接量化為若干離散編碼（Node IDs）。這些 Node IDs 具有以下特點：

緊湊與高效：通常只有 6-15 維，且型別為 int4，可顯著提升推理速度並降低記憶體佔用。
資訊保留充分：在 34 個數據集的節點分類、圖分類、連結預測和帶屬性圖聚類任務中，Node IDs 在高效推理的同時，效能與最先進模型相當或更優。
可解釋性：由於 Node IDs 以離散碼本（codebook）的形式表示多層鄰域結構，易於透過碼字索引理解節點的聚類模式與語義差異。

基於此，我們的發現表明了 GNN 嵌入存在顯著的冗餘，而生成的緊湊離散 Node IDs 提供了圖資料的高階抽象，這一發現可能為圖資料的標記化處理和涉及大模型的應用提供幫助。

方法介紹

我們提出的 NID 框架由兩個階段組成：

生成緊湊的離散 Node IDs：節點透過多層 GNN 編碼，以捕獲多階鄰居結構。在每一層中，節點嵌入被量化為一組結構化的程式碼元組。然後，這些元組被組合成我們所稱之為 Node IDs 的表示。
利用生成的 Node IDs 作為各種下游任務中的節點表示：我們直接使用 Node IDs 進行無監督任務，如節點聚類。我們使用 Node IDs 訓練簡單的 MLP 網路進行監督任務，包括節點分類、連結預測和圖分類。

2.1 Node IDs 的生成

▲ GNN 不同層節點表徵的 t-SNE 視覺化

上圖展示了 GNN 在不同層次生成的節點表示的不同聚類模式。這種多樣性來自於每一層圖卷積的連續應用所造成的平滑效應 [1]。

為了生成結構感知的 Node IDs，我們採用一個層的 GNN 來捕獲多階鄰居結構。在每一層中，我們使用向量量化將 GNN 生成的節點嵌入編碼為個程式碼（整數索引）。對於每個節點，我們將 Node ID 定義為由個程式碼組成的元組，結構如下：

其中表示第層的第個程式碼。

▲ NID 架構圖

如上圖所示，在 GNN 的每一層（）中，我們使用 RVQ [2] 對節點嵌入進行量化，為每個節點生成層次的程式碼。每個程式碼（）是由不同的 codebook生成的，其中是 codebook 的大小。因此，總共有個 codebook，按索引。設表示待量化的向量。

請注意，是 GNN 生成的節點嵌入。當時，表示殘差向量。然後，透過其對應 codebook 中的最近程式碼向量進行近似：

生成程式碼，即最近程式碼向量的索引。

我們提出了一個簡單的通用框架，透過聯合訓練 GNN 和 codebook 來學習 Node IDs（程式碼），並使用以下損失函式：

其中是（自監督/監督）圖學習目標，是向量量化損失。旨在訓練 GNN 生成有效的節點嵌入，而確保 codebook 向量與節點嵌入對齊。對於單個節點，定義為：

其中表示停止梯度操作，是權重引數。上式中的第一項是 codebook loss [3]，它僅影響 codebook，使所選程式碼向量與節點嵌入接近。第二項是 commitment loss [3]，它僅影響節點嵌入，並確保節點嵌入與所選程式碼向量接近。

總結：

我們的 NID 框架與 VQ-VAE [3] 及類似方法（如 VQGraph [4]）在 codebook 學習方面有所不同，我們的訓練目標不涉及使用程式碼向量（）進行重建任務，而是透過圖學習任務（）引導 codebook 學習過程。

近期研究 [5,6] 表明，經過合適調參的經典 GNNs 在節點分類和圖分類任務中仍能達到 SOTA 方法的競爭性表現。因此在實驗中，我們使用了經典 GNNs（GCN、GAT、SAGE、GIN）進行 Node IDs 生成。

2.2 Node IDs 在下游任務中的應用

生成的 Node IDs 可以視為高度緊湊的節點表示，並直接用於各種下游圖學習任務，如下所述。

節點預測任務包括節點分類和節點聚類。對於節點分類，圖中的每個節點都與一個標籤相關聯，表示其類別。我們可以直接利用帶標籤節點的 Node IDs 訓練 MLP 網路進行分類。預測公式為：

對於節點聚類，可以直接將基於向量的聚類演算法，如 -means，應用於 Node IDs 以獲得聚類結果。

邊預測任務通常涉及連結預測。目標是預測任意節點對之間是否應存在一條邊。預測可以透過以下方式進行：

其中表示 Hadamard 積。

圖預測任務包括圖分類和圖迴歸。這些任務涉及為整個圖預測類別標籤或數值。預測可以表達為：

其中對所有 Node IDs 應用全域性均值池化函式，生成圖的表示，然後輸入 MLP 進行預測。注意，選擇讀取函式（如均值池化）被視為一個超引數。

實驗結果

學習得到的 Node IDs 通常由 6-15 個 int4 整陣列成，可作為高效的節點表示。在多個任務中，它們表現出與 SOTA 方法相當或更優的效能，同時顯著提升計算速度和記憶體效率。

3.1 監督表示學習

▲ 監督學習實驗結果

3.2 無監督表示學習

▲ 無監督學習實驗結果

3.2 Node IDs 的分析

▲ Node IDs 的分析

Codebook 崩潰：發現 VQGraph 存在嚴重的 codebook 崩潰問題，即大多數節點被量化為少數幾個程式碼向量，導致大部分 codebook 未被使用。相比之下，NID 實現了高 codebook 使用率，有效避免了 codebook 崩潰。

可解釋性：觀察 Node IDs 的碼字索引分佈，發現其可以有效區分不同類別的節點，具有較強的語義可解釋性。

子圖檢索：利用 Node IDs 的 Hamming 距離進行子圖匹配，效果優於現有的 VQGraph 的結果。

推理加速：在 ogbn-products (節點數百萬級)上，NID 推理時間從 11.9s 降至 0.7ms。

總結

我們提出的 NID 框架可在自監督或有監督訓練過程中，不引入重構損失的前提下，將多層 GNN 的連續嵌入直接量化為緊湊、離散且可解釋的 Node IDs。透過廣泛的實驗與消融研究，我們驗證了 Node IDs 在多工、多資料集上的高效性與實用性。

相比傳統 GNN 嵌入，Node IDs 不僅大幅降低推理時的記憶體與計算開銷，而且在一系列場景中保持了有競爭力甚至更佳的預測效能。這一方法為在大規模圖上實現即時、高效的推理提供了新的見解，也為後續在圖資料與大模型結合（如圖示記化、LLMs 應用）帶來了新的思路。

參考文獻

[1] Qimai Li, Zhichao Han, and Xiao-Ming Wu. Deeper insights into graph convolutional networks for semi-supervised learning. In Thirty-Second AAAI conference on artificial intelligence, 2018.

[2] Julieta Martinez, Holger H Hoos, and James J Little. Stacked quantizers for compositional vector compression. arXiv preprint arXiv:1411.2173, 2014.

[3] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.

[4] Ling Yang, Ye Tian, Minkai Xu, Zhongyi Liu, Shenda Hong, Wei Qu, Wentao Zhang, Bin CUI, Muhan Zhang, and Jure Leskovec. VQGraph: Rethinking graph representation space for bridging GNNs and MLPs. In The Twelfth International Conference on Learning Representations, 2024.

[5] Yuankai Luo, Lei Shi, and Xiao-Ming Wu. Classic gnns are strong baselines: Reassessing gnns for node classification. arXiv preprint arXiv:2406.08993, 2024.

[6] Yuankai Luo, Lei Shi, and Xiao-Ming Wu. Unlocking the potential of classic gnns for graph-level tasks: Simple architectures meet excellence. arXiv preprint arXiv:2502.09263, 2025.

更多閱讀