因果对比学习图神经网络入门指南：从理论到PyTorch实战

1次阅读

共计 3829 个字符，预计需要花费 10 分钟才能阅读完成。

最近在复现 KDD Cup 的医疗诊断赛题时，发现传统 GNN 会把患者的医院访问频率和真实疾病特征高度关联——那些经常来复查的慢性病患者，被模型误判为健康风险更高的人群。这种虚假相关性让我意识到：

特征混淆问题：节点特征和拓扑结构在消息传递过程中相互干扰
过平滑陷阱：多层 GNN 会使不同类别节点的表征趋于相似
因果缺失：模型无法区分相关关系（correlation）和因果关系（causation）

graph LR
    A[原始图] --> B[数据增强模块]
    B --> C[视图 1]
    B --> D[视图 2]
    C --> E[GNN 编码器]
    D --> E
    E --> F[对比损失计算]
    F --> G[梯度回传]

关键组件说明：

视图生成 ：通过边丢弃(node dropping) 和特征掩码 (feature masking) 创建差异化视图
共享编码器：使用 GAT 或 GraphSAGE 作为基础架构，参数在所有视图间共享
投影头：两层的 MLP 将节点表征映射到对比空间

import numpy as np
from torch_geometric.utils import structured_negative_sampling

def generate_pairs(edge_index, num_nodes, walks_per_node=5, walk_length=3):
    """
    基于随机游走的正负样本生成
    :param edge_index: 图的边索引 [2, num_edges]
    :param num_nodes: 节点总数
    :param walks_per_node: 每个节点的游走次数
    :param walk_length: 随机游走长度
    :return: (正样本对, 负样本对)
    """
    pos_pairs = []
    neg_pairs = []

    for _ in range(walks_per_node):
        for node in range(num_nodes):
            current = node
            walk = [current]
            for _ in range(walk_length):
                # 获取邻居节点
                neighbors = edge_index[1, edge_index[0] == current]
                if len(neighbors) == 0:
                    break
                current = np.random.choice(neighbors.numpy())
                walk.append(current)

            # 添加正样本
            for i in range(len(walk)-1):
                pos_pairs.append((walk[i], walk[i+1]))

            # 添加负样本
            for i in range(len(walk)):
                neg_src, neg_dst = structured_negative_sampling(edge_index, num_nodes=num_nodes)
                neg_pairs.append((walk[i], neg_dst.item()))

    return torch.tensor(pos_pairs), torch.tensor(neg_pairs)

对比学习的核心是 InfoNCE 损失：

$$
\mathcal{L} = -\log\frac{\exp(\text{sim}(z_i,z_j)/\tau)}{\sum_{k=1}^N \exp(\text{sim}(z_i,z_k)/\tau)}
$$

其中温度系数 $\tau$ 的调参建议：

小 τ 值（0.05-0.1）：强调困难负样本，适合类别区分度大的场景
大 τ 值（0.5-1.0）：平滑分布，适合细粒度分类任务
自适应 τ ：可尝试根据 epoch 动态调整：
$$\tau = \tau_{\max} – (\tau_{\max}-\tau_{\min})×\frac{\text{current_epoch}}{\text{total_epochs}}$$

import torch
import torch.nn.functional as F
from torch_geometric.nn import GATConv
from torch_geometric.data import Data

class CausalGNN(torch.nn.Module):
    def __init__(self, in_dim, h_dim, out_dim):
        super().__init__()
        self.conv1 = GATConv(in_dim, h_dim, heads=3)
        self.conv2 = GATConv(h_dim*3, h_dim)
        self.projector = torch.nn.Sequential(torch.nn.Linear(h_dim, h_dim),
            torch.nn.ReLU(),
            torch.nn.Linear(h_dim, out_dim)
        )

    def forward(self, x, edge_index):
        x = F.elu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index)  # [num_nodes, h_dim]
        return self.projector(x)  # [num_nodes, out_dim]

def train(model, data, optimizer, tau=0.1):
    model.train()

    # 生成增强视图
    view1 = augment_graph(data)
    view2 = augment_graph(data)

    # 获取对比表征
    z1 = model(view1.x, view1.edge_index)
    z2 = model(view2.x, view2.edge_index)

    # 计算对比损失
    pos_sim = F.cosine_similarity(z1, z2, dim=1) / tau
    neg_sim = torch.mm(z1, z2.t()) / tau  # 矩阵乘法高效计算

    # 对角线是正样本，其余是负样本
    labels = torch.arange(z1.size(0)).to(device)
    loss = F.cross_entropy(neg_sim, labels)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

当 GPU 内存不足时，可以通过梯度累积实现大批量训练：

accum_steps = 4  # 累积 4 个 batch 的梯度

for epoch in range(epochs):
    for i, batch in enumerate(dataloader):
        loss = train_batch(batch)
        loss = loss / accum_steps
        loss.backward()

        if (i+1) % accum_steps == 0:
            optimizer.step()
            optimizer.zero_grad()

邻居采样：使用 Layer-wise 采样避免全图加载

from torch_geometric.loader import NeighborLoader

train_loader = NeighborLoader(
    data,
    num_neighbors=[10, 5],  # 两层采样
    batch_size=512,
    shuffle=True
)

混合精度训练：

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    loss = model(data)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

高度数节点会产生更多负样本，需进行度归一化：

def degree_norm(edge_index, num_nodes):
    row, col = edge_index
    deg = torch.bincount(row, minlength=num_nodes).float()
    norm = 1. / torch.sqrt(deg[row] * deg[col])
    return norm

使用 DDP 时需注意：