BNN不确定性量化在图神经网络中的应用与实现

1次阅读

没有评论

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

图神经网络（GNN）在节点分类、链接预测等任务中表现出色，但传统 GNN 存在一个关键问题：它们通常输出确定性预测，无法量化预测的不确定性。这在以下场景中尤其成问题：

数据包含噪声或缺失值时
测试数据分布与训练数据不同（分布偏移）
需要风险评估的关键应用（如医疗诊断）

传统方法如 Dropout 可以作为简单的 uncertainty 估计，但它们缺乏理论基础，常常低估不确定性。

贝叶斯神经网络（BNN）与传统深度神经网络的关键区别在于权重表示：

传统 DNN：使用确定性的点估计权重
BNN：将权重视为随机变量，学习其概率分布

对于 GNN 而言，这种区别带来了显著优势：

可以自然地从权重分布中采样，得到预测分布
通过蒙特卡洛采样，可以计算预测的方差作为 uncertainty
对模型参数的不确定性有明确建模

在 PyTorch 中，我们可以通过继承 nn.Module 并重写 forward 方法来实现概率权重：

class BayesianLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        # 均值参数
        self.w_mu = nn.Parameter(torch.Tensor(out_features, in_features))
        # 方差参数（实际实现使用 log 方差提高数值稳定性）self.w_rho = nn.Parameter(torch.Tensor(out_features, in_features))
        # 初始化
        nn.init.xavier_normal_(self.w_mu)
        nn.init.constant_(self.w_rho, -3)

    def forward(self, x):
        # 重参数化技巧
        w_epsilon = torch.randn_like(self.w_rho)
        w_sigma = torch.log(1 + torch.exp(self.w_rho))  # softplus
        weight = self.w_mu + w_sigma * w_epsilon
        return F.linear(x, weight)

在预测阶段，我们需要进行多次前向传播（MC 采样）来估计预测分布：

def mc_predict(model, x, num_samples=20):
    outputs = []
    for _ in range(num_samples):
        outputs.append(model(x))
    return torch.stack(outputs)

BNN 需要特殊的损失函数，通常包括两部分：

数据似然项（如交叉熵）
复杂度代价项（KL 散度）

def elbo_loss(output, target, model, kl_weight=0.1):
    # 数据似然项
    likelihood = F.cross_entropy(output, target)

    # 计算所有 BNN 层的 KL 散度和
    kl = 0
    for module in model.modules():
        if isinstance(module, BayesianLinear):
            kl += compute_kl(module.w_mu, module.w_rho)

    return likelihood + kl_weight * kl

下面是一个完整的 Bayesian GNN 实现（基于 PyTorch Geometric）：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class BayesianGCNConv(nn.Module):
    """贝叶斯图卷积层"""
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.w_mu = nn.Parameter(torch.Tensor(out_channels, in_channels))
        self.w_rho = nn.Parameter(torch.Tensor(out_channels, in_channels))
        self.b_mu = nn.Parameter(torch.Tensor(out_channels))
        self.b_rho = nn.Parameter(torch.Tensor(out_channels))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_normal_(self.w_mu)
        nn.init.constant_(self.w_rho, -3)
        nn.init.normal_(self.b_mu, 0, 0.1)
        nn.init.constant_(self.b_rho, -3)

    def forward(self, x, edge_index):
        # 采样权重
        w_sigma = torch.log(1 + torch.exp(self.w_rho))
        w = self.w_mu + w_sigma * torch.randn_like(w_sigma)

        # 采样偏置
        b_sigma = torch.log(1 + torch.exp(self.b_rho))
        b = self.b_mu + b_sigma * torch.randn_like(b_sigma)

        # 执行图卷积
        return F.relu(GCNConv.apply(x, edge_index, w) + b)

class BayesianGNN(nn.Module):
    def __init__(self, num_features, num_classes):
        super().__init__()
        self.conv1 = BayesianGCNConv(num_features, 16)
        self.conv2 = BayesianGCNConv(16, num_classes)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

BNN 的主要计算开销来自：