子图与K-hop图神经网络

概述

标准图神经网络（GNN）采用1-hop消息传递范式——即每次只聚合直接邻居（1跳邻居）的信息。¹这种设计虽然高效，但表达能力受限于1-WL（Weisfeiler-Lehman）测试，无法区分许多非同构图结构。

为了突破这一限制，研究者提出了两大方向：

1. K-hop消息传递：扩展感受野，同时聚合多跳邻居信息
2. 子图GNN：编码局部子图结构，利用子图同构计数提升表达能力

本文系统梳理这两类方法的理论基础、数学 formulation、实现细节及其表达能力分析。

---

一、K-hop消息传递

1.1 核心思想

K-hop消息传递将节点 $v$ 的表示更新从仅依赖1跳邻居扩展到同时依赖 $K$ 跳内的所有邻居：

$${\mathbf{h}}_v^{(l+1)}=\text{UPDATE}\left({\mathbf{h}}_v^{(l)},\left\{{\mathbf{h}}_u^{(l)}:u\in {\mathcal{N}}_G^k(v),1\le k\le K\right\}\right)$$

其中 ${\mathcal{N}}_G^k(v)$ 表示与节点 $v$ 距离为 $k$ 的节点集合（$k$-hop邻居）。

1.2 两种K-hop消息传递范式

NeurIPS 2022论文¹首次严格区分了两种K-hop消息传递机制：

范式一：Multi-hop邻居聚合（MHMA）

$${\mathbf{m}}_v^{(l)}=\text{AGG}\left(\left\{{\mathbf{h}}_u^{(l)}:u\in \bigcup _{k=1}^K{\mathcal{N}}_G^k(v)\right\}\right)$$

将所有 $K$ 跳邻居视为统一集合进行聚合。

范式二：Multi-channel邻居聚合（MCHA）

$${\mathbf{m}}_v^{(l)}=\text{AGG}\left(\left\{{\mathbf{m}}_v^{(l,k)}:1\le k\le K\right\}\right),{\mathbf{m}}_v^{(l,k)}=\text{AGG}\left(\left\{{\mathbf{h}}_u^{(l)}:u\in {\mathcal{N}}_G^k(v)\right\}\right)$$

为每跳单独计算聚合向量，再将 $K$ 个通道的表示融合。

1.3 数学 formulation

定义：设图 $G=(\mathcal{V},\mathcal{E})$，节点 $v$ 的 $K$ 跳邻居集合定义为：

$${\mathcal{N}}_G^K(v)=\bigcup _{k=1}^K\{u\in \mathcal{V}:\text{dist}(u,v)=k\}$$

其中 $\text{dist}(u,v)$ 为最短路径距离。

K-hop消息传递层：

$${\mathbf{h}}_v^{(l+1)}=\sigma \left({\mathbf{W}}^{(l)}\cdot \text{READOUT}\left(\left\{{\mathbf{h}}_v^{(l)}\right\}\cup \left\{{\mathbf{h}}_u^{(l)}:u\in {\mathcal{N}}_G^k(v),k\in [K]\right\}\right)+{\mathbf{b}}^{(l)}\right)$$

1.4 表达能力分析

与1-WL测试的关系

标准1-WL测试在每次迭代中只聚合直接邻居的颜色信息，而K-hop消息传递聚合了多跳邻居的信息。

定理（NeurIPS 2022）：K-hop消息传递的表达能力严格强于1-WL测试，能够区分几乎所有正则图（regular graphs）。¹

然而，研究也揭示了K-hop消息传递的上界：

定理：K-hop消息传递的表达能力被3-WL测试所上界，即：

$$\text{表达能力}(K\text{-hop MP})\le 3\text{-WL}$$

这意味着K-hop消息传递无法区分某些正则图对，例如6节点的两类正则图。

1.5 K-hop vs 更深的GNN

维度	K-hop消息传递	深层1-hop GNN
感受野	单层内覆盖 $K$ 跳	$L$ 层覆盖 $L$ 跳
邻居规模	指数级增长	指数级增长
信息保留	各跳信息可分离	路径信息被压缩
过平滑风险	中等	较高（层数更深）
计算开销	$O(N\cdot d^K)$	$O(N\cdot L\cdot d)$

关键区别在于：深层1-hop GNN通过多层堆叠隐式捕获多跳信息，而K-hop GNN在单层内显式聚合多跳信息，能够保留跳距信息。

---

二、子图神经网络

2.1 动机

子图GNN的核心洞察：节点的表达能力不仅取决于其邻居，还取决于其所在的局部子图结构。

例如，两个度数为3的节点可能位于完全不同的局部结构中：

    A          B
   /|\        /|\
  C D E      F G H
  (三角形)    (路径)

2.2 子图提取方法

Ego-network（自我网络）

以节点 $v$ 为中心，包含 $v$ 及其所有 $k$ 跳内邻居的诱导子图：

$$G_v^k=G[\{v\}\cup {\mathcal{N}}_G^k(v)]$$

$k$-Ego-net

2019年提出的k-hop GNN²使用k-ego-net作为基本计算单元。对于每个节点，提取其k跳自我网络，然后用GNN编码该子图。

2.3 子图同构计数

子图GNN的关键技术之一是子图同构计数——统计某个节点周围包含多少个特定子图模式（如三角形、四边形）。

三角形计数

$$\#△(v)=\frac{1}{2}\sum _{u\in \mathcal{N}(v)}|\mathcal{N}(u)\cap \mathcal{N}(v)|$$

路径计数

$$\#P_3(v)=\left(\genfrac{}{}{0ex}{}{\mathrm{deg}(v)}{2}\right)-\#△(v)$$

2.4 1-WL子图测试

定义：1-WL子图测试对每个节点 $v$ 维护其根植子图（rooted subgraph）的颜色，而不是单节点颜色。

$${\chi }_G^{(t+1)}(v)=\text{HASH}\left({\chi }_G^{(t)}(v),\left\{{\chi }_G^{(t)}(u):u\in G_v^K\right\}\right)$$

定理：1-WL子图测试严格强于1-WL测试，因为子图测试能够感知内部子结构（internal substructure）。³

---

三、KP-GNN框架

3.1 核心思想

NeurIPS 2022论文提出的KP-GNN¹通过注入上下文子结构信息来增强K-hop消息传递的表达能力。

关键洞察：K-hop消息传递采用单向辐射注意力模式（hub-spoke），而1-WL子图GNN采用双向成对注意力模式（pairwise bidirectional）。

K-hop MP:          1-WL Subgraph GNN:
    v                    u1
   /|\                  / | \
  u1 u2 u3             v--v--v
  (中心辐射)           (成对连接)

3.2 子结构编码函数

KP-GNN引入子结构编码函数 $f(\cdot )$ 来编码每个节点的内部子结构：

$$f:G_u^K\times G\to {\mathbb{R}}^d$$

该函数将 $K$ 跳自我子图和全局图作为输入，输出编码内部子结构的 $d$ 维特征。

3.3 多步随机游走编码

论文提出使用多步随机游走的自返回概率来编码子结构：

$$(p_u^1,p_u^2,\dots ,p_u^l)\ne (p_v^1,p_v^2,\dots ,p_v^l)$$

其中 $p_u^k$ 是从节点 $u$ 开始的 $k$ 步随机游走后返回 $u$ 的概率。

定理（随机游走下界）：给定两个 $n$ 节点 $r$ 正则图 $G$ 和 $H$（$3\le r<\sqrt{2\log 2n}$），对于 $K$ 跳自我网络，$2K$ 步随机游走足以区分其内部子结构。³

3.4 SEK-GNN实现

SEK-GNN（Substructure Encoding via K-hop GNN）将子结构编码集成到消息传递框架：

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class SEKGNNLayer(nn.Module):
    """
    SEK-GNN层：结合K-hop消息传递与子结构编码
    """
    def __init__(self, in_channels, out_channels, K=2, 
                 subgraph_encoder=None):
        super().__init__()
        self.K = K
        self.in_channels = in_channels
        self.out_channels = out_channels
        
        # 节点特征变换
        self.lin_self = nn.Linear(in_channels, out_channels)
        self.lin_neigh = nn.Linear(in_channels, out_channels)
        
        # 多跳聚合权重
        self.hop_weights = nn.Parameter(torch.ones(K) / K)
        
        # 子结构编码器
        self.subgraph_encoder = subgraph_encoder or nn.Identity()
        
        # 融合层
        self.fusion = nn.Linear(out_channels * 2, out_channels)
    
    def forward(self, x, edge_index, subgraph_features):
        """
        参数:
            x: 节点特征 (N, in_channels)
            edge_index: 边索引 (2, E)
            subgraph_features: 子图编码特征 (N, d_sub)
        """
        N = x.shape[0]
        
        # 1. 自环特征
        h_self = self.lin_self(x)
        
        # 2. K-hop邻居聚合
        h_multi_hop = self.compute_khop_aggregation(x, edge_index)
        
        # 3. 子结构特征编码
        h_sub = self.subgraph_encoder(subgraph_features)
        
        # 4. 特征融合
        h_fused = torch.cat([h_self + h_multi_hop, h_sub], dim=-1)
        h_out = self.fusion(h_fused)
        
        return F.relu(h_out)
    
    def compute_khop_aggregation(self, x, edge_index):
        """计算多跳邻居加权聚合"""
        N = x.shape[0]
        h_agg = torch.zeros(N, self.out_channels, device=x.device)
        
        # 获取邻居列表
        row, col = edge_index
        
        # 对每跳加权聚合
        hop_weights = F.softmax(self.hop_weights, dim=0)
        
        for k in range(self.K):
            # 找到k+1跳邻居（简化实现）
            neighbors_k = col  # 实际需根据距离筛选
            
            if len(neighbors_k) > 0:
                h_neigh = self.lin_neigh(x[neighbors_k])
                # 按目标节点聚合
                h_agg = h_agg + hop_weights[k] * torch.zeros(N, self.out_channels, device=x.device).scatter_add(
                    0, neighbors_k.unsqueeze(-1).expand(-1, self.out_channels), h_neigh)
        
        return h_agg

3.5 表达能力保证

定理（SEK-1-WL表达能力）：SEK-1-WL测试：

1. 严格强于 1-WL子图测试
2. 严格强于 K-hop 1-WL测试
3. 表达能力不低于 3-WL测试

这意味着SEK-GNN能够区分标准K-hop GNN和标准子图GNN无法区分的图结构。

---

四、表达能力对比

4.1 WL测试层次结构

$$\text{1-WL}<\text{K-hop WL}<\text{1-WL Subgraph}\le \text{SEK-1-WL}\le \text{3-WL}<\text{2-WL}<\cdots$$

4.2 图同构测试对比

测试方法	表达能力	计算复杂度	可扩展性
1-WL	基础	$O(N)$	优秀
K-hop WL	较强	$O(NK)$	良好
1-WL Subgraph	强	$O(N \cdot	\text{subgraph}
SEK-1-WL	很强	$O(NK+\text{RW steps})$	良好
3-WL	非常强	$O(N^3)$	差
k-WL	极强	$O(N^k)$	极差

4.3 能区分 vs 不能区分的图

图类型	1-WL	K-hop	Subgraph	SEK
链 vs 环	✗	✓	✓	✓
三角形 vs 四边形	✓	✓	✓	✓
同构正则图	✗	部分	✓	✓
距离正则图	✗	✗	✗	✓

---

五、实际实现

5.1 K-hop邻居提取

import torch
from torch_geometric.utils import k_hop_subgraph
 
def extract_khop_neighbors(edge_index, node_idx, num_nodes, K=2):
    """
    提取节点node_idx的K跳邻居
    
    返回:
        subset: K跳子图中的节点列表
        edge_index_sub: K跳子图的边
        mapping: 原始节点到子图节点的映射
        edge_mask: 边掩码
    """
    subset, edge_index_sub, mapping, edge_mask = k_hop_subgraph(
        node_idx, K, edge_index, relabel_nodes=True, num_nodes=num_nodes
    )
    return subset, edge_index_sub, mapping, edge_mask
 
def batch_extract_khop_features(x, edge_index, K=2):
    """
    批量提取K-hop特征
    
    参数:
        x: 节点特征 (N, d)
        edge_index: 边索引 (2, E)
        K: 跳数
    
    返回:
        hop_features: 每跳的聚合特征列表
    """
    N = x.shape[0]
    hop_features = []
    
    for k in range(1, K + 1):
        # 计算k跳邻居掩码
        k_hop_agg = torch.zeros_like(x)
        
        # 使用PyG的k_hop_subgraph逐节点提取
        for v in range(N):
            subset, _, _, _ = k_hop_subgraph(v, k, edge_index, num_nodes=N)
            # 聚合k跳邻居（不含中心节点）
            k_hop_neighbors = subset[subset != v]
            if len(k_hop_neighbors) > 0:
                k_hop_agg[v] = x[k_hop_neighbors].mean(dim=0)
        
        hop_features.append(k_hop_agg)
    
    return hop_features

5.2 完整K-hop GNN模型

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import k_hop_subgraph
 
class KHopGNNConv(MessagePassing):
    """
    K-hop消息传递卷积层
    """
    def __init__(self, in_channels, out_channels, K=2, aggr='mean'):
        super().__init__(aggr=aggr)
        self.K = K
        self.lin = nn.Linear(in_channels, out_channels)
    
    def forward(self, x, edge_index):
        """
        参数:
            x: 节点特征 (N, in_channels)
            edge_index: 边索引 (2, E)
        """
        # 初始化多跳特征
        h_multihop = torch.zeros_like(x)
        h_multihop = self.lin(x)
        
        # 对每跳分别聚合
        hop_weights = torch.ones(self.K) / self.K
        h_agg = torch.zeros_like(x)
        
        for k in range(1, self.K + 1):
            # 获取k跳邻居子图
            subset, edge_index_k, _, _ = k_hop_subgraph(
                torch.arange(len(x)), k, edge_index, num_nodes=len(x)
            )
            
            # 在子图上执行消息传递
            h_k = self.propagate(edge_index_k, x=self.lin(x))
            h_agg = h_agg + hop_weights[k-1] * h_k
        
        return F.relu(h_multihop + h_agg)
    
    def message(self, x_j):
        """消息函数"""
        return x_j
 
 
class KHopGNN(nn.Module):
    """
    K-hop GNN模型
    """
    def __init__(self, in_channels, hidden_channels, out_channels, K=2, num_layers=3):
        super().__init__()
        self.K = K
        
        self.convs = nn.ModuleList()
        self.convs.append(KHopGNNConv(in_channels, hidden_channels, K))
        
        for _ in range(num_layers - 2):
            self.convs.append(KHopGNNConv(hidden_channels, hidden_channels, K))
        
        self.convs.append(KHopGNNConv(hidden_channels, out_channels, K))
        
        self.dropout = nn.Dropout(0.5)
    
    def forward(self, x, edge_index):
        for i, conv in enumerate(self.convs):
            x = conv(x, edge_index)
            if i < len(self.convs) - 1:
                x = F.relu(x)
                x = self.dropout(x)
        
        return x

5.3 子结构编码实现

import torch
import torch.nn as nn
import random
 
class RandomWalkSubgraphEncoder(nn.Module):
    """
    基于随机游走的子结构编码器
    用于KP-GNN/SEK-GNN的子结构信息注入
    """
    def __init__(self, num_nodes, embedding_dim, walk_steps=4):
        super().__init__()
        self.walk_steps = walk_steps
        
        # 节点嵌入
        self.node_embedding = nn.Embedding(num_nodes, embedding_dim)
        
        # 随机游走概率（可学习）
        self.transition_matrix = None  # 延迟初始化
        
        # 多步自返回概率编码
        self.return_encoder = nn.Linear(walk_steps, embedding_dim)
    
    def compute_return_probability(self, edge_index, num_nodes, start_node, steps):
        """
        计算从start_node开始的随机游走在steps步后返回的概率
        
        参数:
            edge_index: 边索引
            num_nodes: 节点数
            start_node: 起始节点
            steps: 步数
        
        返回:
            return_prob: 自返回概率
        """
        # 构建邻接表
        adj = torch.zeros(num_nodes, num_nodes)
        adj[edge_index[0], edge_index[1]] = 1
        adj[edge_index[1], edge_index[0]] = 1  # 无向图
        
        # 行归一化得到转移矩阵
        degree = adj.sum(dim=1, keepdim=True)
        transition = adj / (degree + 1e-8)
        
        # 随机游走模拟
        current = start_node
        for _ in range(steps):
            probs = transition[current]
            current = torch.multinomial(probs + 1e-8, 1).item()
        
        # 返回概率
        return_prob = 1.0 if current == start_node else 0.0
        return return_prob
    
    def forward(self, x, edge_index, node_indices):
        """
        参数:
            x: 节点特征
            edge_index: 边索引
            node_indices: 需要编码的节点索引
        
        返回:
            subgraph_features: 子结构特征
        """
        num_nodes = x.shape[0]
        batch_size = len(node_indices)
        
        # 计算多步自返回概率
        return_probs = torch.zeros(batch_size, self.walk_steps)
        
        for i, v in enumerate(node_indices):
            for step in range(1, self.walk_steps + 1):
                # 多次采样估计概率
                probs = []
                for _ in range(10):  # 采样次数
                    prob = self.compute_return_probability(
                        edge_index, num_nodes, v, step
                    )
                    probs.append(prob)
                return_probs[i, step - 1] = sum(probs) / len(probs)
        
        # 编码自返回概率
        h_sub = self.return_encoder(return_probs)
        
        # 结合节点嵌入
        node_emb = self.node_embedding(node_indices)
        
        return h_sub + node_emb

5.4 内存与计算考量

操作	时间复杂度	空间复杂度	优化策略
K-hop邻居提取	$O(N\cdot d^K)$	$O(N+E)$	预计算 + 缓存
随机游走编码	$O(N\cdot L\cdot R)$	$O(N)$	向量化采样
子图聚合	$O(N \cdot	S	)$

其中 $N$ 为节点数，$K$ 为跳数，$d$ 为特征维度，$L$ 为随机游走步数，$R$ 为采样次数，$|S|$ 为子图大小。

---

六、应用场景与实践建议

6.1 何时使用K-hop GNN

适合场景：

1. 需要捕获远程依赖：如分子中原子间的远程相互作用
2. 节点角色需要多跳上下文：如社交网络中识别社区结构
3. 深层GNN效果不佳：当过平滑严重时，K-hop可减少层数
4. 图直径较小：K-hop能够覆盖整个图时效果最佳

不适合场景：

1. 大规模图：K-hop邻居数量指数增长
2. 局部特征足够：只需要直接邻居信息
3. 实时推理：K-hop计算开销较大

6.2 何时使用子图GNN

适合场景：

1. 结构识别任务：如分子性质预测（需要识别官能团）
2. 角色区分：识别节点在局部结构中的角色
3. 正则图区分：需要区分同构的正则图结构

6.3 K与子图大小的选择

K值/子图大小	表达能力	计算开销	建议
K=1	基础	低	默认选择
K=2	中等	中等	平衡选择
K=3	较强	较高	小图使用
全子图	最强	最高	小规模图

6.4 与其他方法的结合

1. K-hop + 注意力：使用注意力加权不同跳的信息
2. 子图 + 位置编码：结合绝对位置与相对结构
3. K-hop + 残差连接：缓解过平滑问题
4. SEK-GNN + 图transformer：融合子结构与全局注意力

---

七、相关工作与扩展

7.1 近期研究进展

年份	方法	核心贡献
2019	k-hop GNN2	首次提出K-hop消息传递框架
2020	DIGL4	扩散图学习，结合图扩散核
2022	KP-GNN1	分析K-hop表达能力，引入KP-GNN框架
2024	SEK-GNN3	子结构编码，超越3-WL表达能力

7.2 相关技术

• Nested GNN：对每个节点应用子GNN编码其邻居子图
• ID-GNN：标记根节点，增强自我网络表达能力
• GNN-As-Kernel：使用GNN作为核函数处理子图
• DS-GNN：方向性子图GNN

详见 图神经网络。

---

八、总结

K-hop消息传递和子图GNN代表了增强GNN表达能力的两条主要路径：

1. K-hop消息传递通过扩展感受野，在单层内捕获多跳信息，表达能力介于1-WL和3-WL之间
2. 子图GNN通过编码局部子图结构，能够感知内部子结构，表达能力强于K-hop消息传递
3. KP-GNN/SEK-GNN通过注入子结构信息，结合两者优势，可达到3-WL级别的表达能力

在实际应用中，应根据任务需求、图规模和计算资源选择合适的方法。

---

参考

---

Footnotes

1. Feng et al., “How Powerful are K-hop Message Passing Graph Neural Networks”, NeurIPS 2022, https://arxiv.org/abs/2205.13328 ↩ ↩² ↩³ ↩⁴ ↩⁵

2. Morris et al., “k-hop Graph Neural Networks”, arXiv 2019, https://arxiv.org/abs/1907.06051 ↩ ↩²

3. Feng et al., “Improving the Expressiveness of K-hop Message-Passing GNNs by Injecting Contextualized Substructure Information”, KDD 2024, https://arxiv.org/abs/2406.19244 ↩ ↩² ↩³

4. Haseleu et al., “Diffusion-based Graph Neural Networks”, arXiv 2020 ↩