PyTorch中GNN的Integrated Gradients扩展：获取单节点分数

我正在为自研的图神经网络（GNN）实现Integrated Gradients计算，参考论文《Axiomatic Attribution for Deep Networks》中的公式3。网络输入为N×M矩阵，代表含N个顶点的图，每个顶点对应M维特征向量。论文方法针对N维输入推导，我希望为每个节点生成单独分数，当前实现得到N×M矩阵后通过torch.sum(dim=1)压缩，但不够简洁。

graph_copy = input_graph.detach().clone()
# k/m in the formula
k_m = torch.linspace(0, 1, n_steps)
# Array to store the samples
input_features = [baseline.clone()]    # The first sample is the baseline
# x-x'
diff = original_input_feature - baseline

# Populates the samples array
for i in range(1, k_m.shape[0]):
    temp = baseline + k_m[i] * diff
    input_features_path.append(temp)

gradients = []
for i in range(k_m.shape[0]):
    ### Zeroing the gradient of the vectors for each input 
    input_features[i].requires_grad = True
    model.zero_grad()
    task.zero_grad()
    
    temp_model_out = model(graph = graph_copy, input = input_features[i])['graph_feature'][0]
    temp_mlp_output = task.mlp(temp_model_out)
    temp_prob = Softmax(temp_mlp_output, dim = 0)

    temp_gradient = grad(outputs = temp_prob[true_label_id], inputs = input_features_path[i])
    gradients.append(temp_gradient[0])

    input_features_path[i].requires_grad = False

with torch.no_grad():
    ig_scores = original_input_feature * torch.stack(gradients, dim = 2).mean(dim = 2)    
    final_ig_scores = ig_scores.sum(dim = 1)
    sorted, indices = torch.sort(final_ig_scores, descending = True)

通过向量化操作替代循环，同时整合节点分数的求和逻辑，让代码更简洁紧凑：

1. 批量生成插值样本
   用广播机制一次性生成所有路径样本，避免循环拼接：

# 生成所有插值样本，shape为(n_steps, N, M)
input_features_path = baseline + k_m[:, None, None] * diff[None, :, :]
# 确保第一个样本为基线（和原逻辑对齐）
input_features_path[0] = baseline.clone()

2. 批量计算梯度
   使用PyTorch 2.0+的torch.func.vmap实现梯度批量计算，消除循环：

from torch.func import vmap, grad

def compute_gradient(sample):
    sample.requires_grad_(True)
    model.zero_grad()
    task.zero_grad()
    model_out = model(graph=graph_copy, input=sample)['graph_feature'][0]
    mlp_out = task.mlp(model_out)
    prob = torch.softmax(mlp_out, dim=0)
    return grad(lambda x: prob[true_label_id])(sample)

# 批量计算所有样本的梯度，shape为(n_steps, N, M)
gradients = vmap(compute_gradient)(input_features_path)

3. 简化节点分数计算
   直接在IG计算后沿特征维度求和，省去中间变量：

with torch.no_grad():
    # 一步得到节点级IG分数，shape为(N,)
    final_ig_scores = (original_input_feature * gradients.mean(dim=0)).sum(dim=1)
    sorted_scores, indices = torch.sort(final_ig_scores, descending=True)

上述优化既保留了原逻辑的正确性，又消除了冗余循环，同时将节点分数的求和逻辑整合到最终计算步骤中，代码结构更简洁。

内容的提问来源于stack exchange，提问作者miguel pedraza