Stanford CS224W: Machine Learning with Graphs

By Ivan Ge and Flora Yuan as part of the Stanford CS 224W Final Project


代码下载：https://t.zsxq.com/iH0Gn

请索引第42个项目


Quantum Graph Neural Networks Applied | by Pavan Jayasinha | Medium


图神经网络（GNN）近年来作为一种强大的深度学习模型脱颖而出，其设计目的是处理图结构数据，其中实体（节点）之间通过关系（边）相互连接[1]。与处理固定空间（如网格或序列）数据的传统神经网络（例如Transformer或卷积神经网络）不同，GNN通过学习底层拓扑结构来处理不规则结构。本文将探讨如何利用图神经网络进行高能物理实验中的异常检测。


近年来，高能物理研究的一大重点是利用快速机器学习进行异常检测[2]。诸如ATLAS之类的实验通过分析大型强子对撞机（Large Hadron Collider，LHC）产生的粒子碰撞数据，研究了广泛的粒子物理现象。然而，LHC每天产生的数据量高达数十亿个事件（过滤后约为1PB）[3]。绝大多数事件都由人们熟知的标准模型过程（我们称之为背景信号）构成，而只有极少数可能包含新物理的痕迹（我们称之为新信号）。为了解决海量数据解析的难题，研究人员依赖于无监督机器学习方法，例如异常检测，来帮助标记并存储可能代表偏离标准模型的数据。


传统的异常检测方法通常将碰撞事件视为扁平的特征向量，这可能会忽略粒子间复杂的相互作用关系。每个记录到的碰撞事件都可能包含多个粒子（电子、μ子、喷注、缺失横向能量），这些粒子之间存在着空间和运动学上的关系，而这些关系对于理解事件结构至关重要。在本篇博文中，我们将利用图神经网络（GNN）构建一个基于GNN的异常检测模型来解决这些问题。


问题


标准模型无法解释我们对暗物质和暗能量的观测结果，这意味着它并不完整。也就是说，宇宙中存在着我们尚未发现的“新”物理现象！为了寻找标准模型之外的物理现象，我们开展粒子物理实验并分析粒子碰撞数据（就像2012年发现希格斯玻色子那样！）[4]。识别标准模型无法解释的新事件对于构建更完整的宇宙模型至关重要。


本文仅关注这些实验的数据方面：如何只存储有价值的数据并舍弃所有其他数据？由于没有标签，我们的主要方法是图级无监督异常检测。每个碰撞事件都被编码成一个图，然后使用我们的模型预测该图是背景信号还是新信号。通过使用包含数百万个背景数据点（即标准模型交互）的数据集和新信号测试集（不属于标准模型的交互），我们希望训练一个能够区分这些信号的模型。


数据集


我们使用此处提供的数据集。数据集中的每个数据点代表一次观测到的碰撞事件，并包含一系列测量得到的粒子属性。对于每个事件，我们拥有以下信息：


  • 缺失横向能量（MET）：两个值描述了碰撞中似乎“消失”的能量（即，碰撞后未从观测到的粒子中得到的能量）。

  • 电子：事件中最多可探测到四个电子（e0 至 e3），每个电子都由三个数字描述，分别表示其动量和方向。这些数字分别用 eta、p_t 和 phi 表示。

  • μ子：类似地，最多四个μ子（μ0到μ3），每个μ子都有相同的三个测量值。

  • 喷流：最多十个喷流（粒子簇，j0 到 j9），每个喷流同样由三个运动学特征表示。


第一个数据集包含背景事件，这些事件是我们在常规标准模型物理学中预期发生的标准粒子碰撞。我们有数百万个这样的事件，它们代表了正常对撞机数据的样貌。


接下来的四个数据集包含一些根据我们对标准模型的理解而言罕见且出乎意料的新信号事件。这里使用的四个示例是：


  • A → 4ℓ：一种衰变为四个轻子的奇异粒子

  • H → ττ：希格斯玻色子衰变为一对τ粒子

  • H → τν：带电希格斯粒子衰变成 tau 粒子和中微子

  • 轻子夸克产生：连接轻子和夸克的假想粒子


数据处理


我们首先将原始数值事件数据转换为图结构，以便我们的图神经网络（GNN）模型能够理解它。每个事件都成为一个图，其中节点代表单个粒子的特征（例如，一个节点可能是 e0_pt，即测量到的第一个电子的横向动量），边则表示这些特征之间的关系。下面，我们将探讨定义这些边的几种方法，每种方法都旨在捕捉底层物理的不同方面。当然，还有许多其他方法可以探索这方面的内容！


1. 链边（基线法）


最简单的方法是按固定的顺序连接各个要素——本质上是形成一条贯穿所有节点的长链。


2. ΔR 边缘


在高能物理学中，角距离可以表示为



它通常用于测量两个粒子在探测器空间中的接近程度。我们根据阈值，将粒子间距离 ΔR 值接近的节点连接起来。


3. 能量关联边缘


高能粒子之间常常相互影响或关联。本文中，我们基于节点归一化横向动量 (pₜ) 的相似程度连接节点，并同样使用可调阈值。


4. k近邻（kNN）边


我们不直接运用物理原理，而是着眼于特征空间。对于每个节点，我们根据特征向量的数值相似性，将其与其k 个最近邻节点连接起来。



特征子集（56 个节点中的前 16 个节点）的示例图


模型架构——图自动编码器


我们使用图自编码器（GAE）来学习正常事件的结构。本质上，我们仅使用背景事件训练模型，使模型能够很好地重构背景信号的图。目标是，对于新信号，信号事件的重构效果会很差，与原始输入图相比会产生较大的误差。


模型架构图


编码器：编码器使用堆叠的 GNN 卷积层，使每个节点能够从其邻居节点收集信息。经过多次使用 GraphConv 层的消息传递步骤后，整个图被压缩成一个紧凑的潜在向量，该向量表示模型潜在空间中输入图的概要信息。


我们尝试了以下几种图神经网络架构：


图卷积网络（GCN）：


图卷积网络（GCN）[5] 是一种专为图结构数据设计的卷积神经网络（CNN），它通过图卷积运算聚合相邻节点的特征来计算节点嵌入。在我们的实现中，编码器使用 GCN 层来学习图表示，以捕捉节点特征之间的关系。它通过应用图卷积来处理每个事件图，从而聚合来自相邻节点的信息。然后，使用均值池化对学习到的节点嵌入进行全局池化，生成图级潜在表示，解码器随后对其进行解码。


GCN解码器随后利用学习到的图级潜在表示来重构原始节点特征。具体做法是将潜在向量扩展到每个节点，然后应用图卷积将信息反向传播到整个图结构中。解码器学习重构与给定输入相匹配的节点特征。其目标是使正常事件的重构误差较低，而异常事件的重构误差较高。


我们的 PyG [6] 实现如下所示：

classGCNAnomaly(nn.Module):
def__init__(self, in_channels, hidden_channels, latent_dim, num_layers):
super().__init__()
self.encoder_convs = nn.ModuleList()

# GCN layers for the encoder
self.encoder_convs.append(GCNConv(in_channels, hidden_channels))
for _ inrange(num_layers - 2):
self.encoder_convs.append(GCNConv(hidden_channels, hidden_channels))
self.encoder_convs.append(GCNConv(hidden_channels, hidden_channels))
self.encoder_proj = nn.Linear(hidden_channels, latent_dim)

# GCN layers for the decoder
self.decoder_proj = nn.Linear(latent_dim, hidden_channels)
self.decoder_convs = nn.ModuleList()
for _ inrange(num_layers - 2):
self.decoder_convs.append(GCNConv(hidden_channels, hidden_channels))
self.decoder_convs.append(GCNConv(hidden_channels, in_channels))

defencode(self, x, edge_index, batch):
for conv inself.encoder_convs:
            x = conv(x, edge_index).relu()
        g = global_mean_pool(x, batch)
returnself.encoder_proj(g)

defdecode(self, z, edge_index, batch):
        z_expanded = z[batch]
# expand back into graph dimensions before convolutions
        x = self.decoder_proj(z_expanded).relu()
for conv inself.decoder_convs[:-1]:
            x = conv(x, edge_index).relu()
returnself.decoder_convs[-1](x, edge_index)


GraphSAGE：


GraphSAGE [7] 是另一种类型的图神经网络 (GNN)，它也通过聚合邻域特征信息来学习节点表示。然而，与使用固定卷积滤波器的图卷积网络 (GCN) 不同，GraphSAGE 使用可学习的聚合函数（例如均值、最大值或 LSTM）来组合邻域特征，使其成为一个更灵活的模型。


我们采用 GraphSAGE 架构作为编码器，但为了测试 GNN 类型在利用编码器创建准确潜在空间方面的差异，我们对所有模型都采用了 GCN 解码器。


我们同样使用 GraphSAGE 进行异常检测，通过训练我们的模型，以类似于我们用于 GCN 的方法，从图级表示中重建节点特征。


以下是我们用 PyG 实现的 GraphSAGE：

classGraphSAGEAnomaly(nn.Module):
def__init__(self, in_channels, hidden_channels, latent_dim, num_layers, sage_agg="mean"):
super().__init__()
self.encoder_convs = nn.ModuleList()

# GraphSAGE layers for the encoder
self.encoder_convs.append(SAGEConv(in_channels, hidden_channels, aggr=sage_agg))
for _ inrange(num_layers - 2):
self.encoder_convs.append(SAGEConv(hidden_channels, hidden_channels, aggr=sage_agg))
self.encoder_convs.append(SAGEConv(hidden_channels, hidden_channels, aggr=sage_agg))
self.encoder_proj = nn.Linear(hidden_channels, latent_dim)

# Decoder uses GCN
self.decoder_proj = nn.Linear(latent_dim, hidden_channels)
self.decoder_convs = nn.ModuleList()
for _ inrange(num_layers - 2):
self.decoder_convs.append(GCNConv(hidden_channels, hidden_channels))
self.decoder_convs.append(GCNConv(hidden_channels, in_channels))


GIN（图同构网络）：


GIN [8] 是另一种图神经网络架构，其设计目标是与 Weisfeiler-Lehman 图同构测试一样强大，因此在区分非同构图方面尤为有效。为此，它在每个图卷积层中使用多层感知器 (MLP) 来转换节点特征，使其能够捕获比标准 GCN 更复杂的图结构。


为了利用 GIN 架构进行异常检测，我们的模型首先通过 GIN 层处理图，这些层对节点特征应用 MLP 变换，同时聚合邻居信息。这种基于 MLP 的消息传递机制使模型能够学习更复杂的特征变换，从而更好地区分输入图中的细微差别。最后，我们再次使用全局池化层来形成图级表示，该表示编码了整体事件结构。


然后，我们利用图卷积网络（GCN）解码器从编码器生成的潜在表示中重构节点特征，并将重构误差作为异常得分。该方法背后的动机在于，专门关注非同构图测试可能有利于高能物理中的异常检测，因为在高能物理中，细微的差异即可区分背景信号和新信号。


我们用 PyG 实现的 GIN：

classGINAnomaly(nn.Module):
def__init__(self, in_channels, hidden_channels, latent_dim, num_layers):
super().__init__()
self.encoder_convs = nn.ModuleList()

# MLP layer, consisting of two linear layers w/ ReLU
        mlp = nn.Sequential(
            nn.Linear(in_channels, hidden_channels),
            nn.ReLU(),
            nn.Linear(hidden_channels, hidden_channels)
        )
# GIN conv layers
self.encoder_convs.append(GINConv(mlp, train_eps=True))
for _ inrange(num_layers - 2):
            mlp = nn.Sequential(
                nn.Linear(hidden_channels, hidden_channels),
                nn.ReLU(),
                nn.Linear(hidden_channels, hidden_channels)
            )
self.encoder_convs.append(GINConv(mlp, train_eps=True))
self.encoder_proj = nn.Linear(hidden_channels, latent_dim)

# Decoder to reconstruct node features (uses GCN)
self.decoder_proj = nn.Linear(latent_dim, hidden_channels)
self.decoder_convs = nn.ModuleList()
for _ inrange(num_layers - 2):
self.decoder_convs.append(GCNConv(hidden_channels, hidden_channels))
self.decoder_convs.append(GCNConv(hidden_channels, in_channels))


GAT（图注意力网络）


GAT架构[9]将之前在Transformer模型中使用的注意力机制引入到图神经网络中，使得节点能够关注具有不同重要性权重的不同邻居。这使得模型在聚合每个节点的邻域信息时，能够专注于最相关的连接。


在我们的异常检测框架中，我们使用GAT架构作为编码器，以帮助识别对理解事件结构至关重要的重要粒子相互作用。注意力机制会计算每个相邻边的权重。这在粒子物理学中可能特别有用，因为某些粒子相互作用可能比其他相互作用更能指示异常过程。


GAT编码器采用多头注意力机制来计算邻域特征的加权聚合，每个注意力头关注邻域的不同方面。这些学习到的注意力权重通过突出显示模型认为重要的边，从而增强了模型的可解释性，这有助于物理学家理解模型使用哪些特征来区分异常。与其他模型一样，学习到的节点嵌入随后被全局池化，以形成事件的图级表示。


我们用 PyG 实现的 GAT：

classGATAnomaly(nn.Module):
def__init__(self, in_channels, hidden_channels, latent_dim, num_layers, dropout=0.1):
super().__init__()
self.encoder_convs = nn.ModuleList()

# GAT layers for the encoder
self.encoder_convs.append(GATConv(in_channels, hidden_channels, heads=1, concat=False, dropout=dropout))
for _ inrange(num_layers - 2):
self.encoder_convs.append(GATConv(hidden_channels, hidden_channels, heads=1, concat=False, dropout=dropout))
self.encoder_convs.append(GATConv(hidden_channels, hidden_channels, heads=1, concat=False, dropout=dropout))
self.encoder_proj = nn.Linear(hidden_channels, latent_dim)

# Decoder that reconstructs node features (GCN)
self.decoder_proj = nn.Linear(latent_dim, hidden_channels)
self.decoder_convs = nn.ModuleList()
for _ inrange(num_layers - 2):
self.decoder_convs.append(GCNConv(hidden_channels, hidden_channels))
self.decoder_convs.append(GCNConv(hidden_channels, in_channels))


解码器：解码器从潜在表示中重构节点特征。它获取图级潜在向量（来自全局均值池化），并将其扩展到图中的每个节点。然后，解码器应用一系列图卷积网络（GCN）层，将信息传播到图结构中，从而重构原始节点特征。重构过程首先将潜在向量投影到隐藏维度，然后应用图卷积，利用编码过程中使用的相同边结构。最后一层输出与输入维度相同的重构节点特征。


原始节点特征与重建节点特征之间的均方误差用作异常得分。理想情况下，与训练过程中学习到的事件相匹配的正常背景事件能够实现较低的重建误差，而新信号数据中的异常事件则会产生较高的误差，表明它们偏离了模型预期的标准背景分布。

"""Code to decode latent to reconstructed node features."""

# Num of nodes in each batch
if batch isNone:
            num_nodes = edge_index.max().item() + 1if edge_index.numel() > 0else1
            batch = z.new_zeros(num_nodes, dtype=torch.long)
else:
            num_nodes = batch.size(0)

# Expand latent to each node (one latent per graph)
# z shape: [batch_size, latent_dim]
        z_expanded = z[batch]  # [num_nodes, latent_dim]

# Project to hidden layer
        x = self.decoder_proj(z_expanded)
        x = F.relu(x)

# GCN layers in the decoder
for conv inself.decoder_convs[:-1]:
            x = conv(x, edge_index)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)

# final reconstruction layer
        x = self.decoder_convs[-1](x, edge_index)
return x


训练


为了训练我们的模型，我们使用了以下损失函数和超参数：


  • 损失函数：原始节点特征与重建节点特征之间的均方误差

  • 优化器：Adam，学习率为 1e-3

  • 正则化：Dropout（p = 0.1）和 L2 正则化

  • 所用的阈值参数为 ΔR = 0.5，pₜ = 0.1，k = 5。


对于训练集/验证集/测试集的划分，我们使用90/10的比例，将前380万个数据点划分为训练集和验证集，用于训练模型。然后，我们使用最后20万个事件进行测试。


评估


自编码器在背景数据上训练完成后，我们使用预留的背景数据测试集（约20万个数据点）以及新的信号数据对其进行评估。对于每个事件，我们计算重构误差，即解码器生成的重构节点特征与原始节点特征之间的均方误差。然后，我们绘制真阳性率 (TPR) 与假阳性率 (FPR) 的 ROC 曲线，并计算曲线下面积 (AUC)。AUC 值越高，表示背景和信号之间的分离效果越好。此外，我们还使用 TPR @ FPR = 1e-5 指标。该指标量化了在极低 FPR 下的真阳性率，旨在确保将最少的假阳性结果保留下来，用于实际生产模型的后续数据分析。


结果


我们提供了一些不同模型的 ROC-AUC 曲线示例，以及 FPR = 1e-5 时的 TPR。更完整的测试结果可在 GitHub 上找到。




4 种不同模型架构的 ROC-AUC 曲线示例


我们对不同的模型架构类型和边类型进行了测试，并展示了每种不同的新信号数据集的结果。


轻子夸克：


从 ROC-AUC 结果可以看出，总体而言，这些模型在区分背景和轻子夸克产生这一新信号事件方面表现良好。在此指标下，GIN 和 GraphSAGE 模型表现最佳。然而，对于极端背景抑制，我们发现大多数模型的真阳性率 (TPR) 为 0，只有基于能量的 GCN 模型和基于 ΔR 的 GraphSAGE 模型例外。


从这些数据中我们可以得出以下几点结论：


  • 高 ROC-AUC 值并不意味着在极低的假阳性率 (FPR) 下也有效，而极低的假阳性率对于高能物理领域的异常检测至关重要，因为该领域每天都会产生数十亿个事件。

  • 我们发现，与其他模型相比，GraphSAGE 显著提高了尾部灵敏度，尽管大多数轻夸克事件在极低的 FPR 下仍然没有表现出异常。


H 到 Tau Nu


从 ROC-AUC 值可以看出，所有模型的异常检测性能总体上都很出色。基于能量边的 GCN 模型在该指标上的表现最差，但有趣的是，它在极端背景抑制下进行信号恢复时表现相当不错。与之前一样，我们发现基于 ΔR 边的 GraphSAGE 模型仍然比其他模型性能更优。


H 到 Tau Tau


我们看到，此处的 ROC-AUC 值比之前的 H-to-tau nu 信号弱，但仍与背景信号有一定的区分度。我们注意到，除基于 ΔR 边的 GraphSAGE 模型外，所有模型的 TPR 值均非零，但都非常小。然而，所有模型的 TPR 值仍然不为零。


A 至 4 l


有趣的是，我们注意到几乎所有模型的真阳性率 (TPR) 都相同。这并不完全出乎意料，因为在极高的假阳性率 (FPR) 水平下，异常值阈值（重建损失）设置得非常高，这意味着只有最极端的异常值才会被保留。这表明，不同模型的异常值分布在其极端尾部非常相似。


总体要点


我们的实验表明，图自编码器有望成为高能物理（HEP）异常检测中一种可行的无监督方法[10]。然而，其有效性似乎取决于信号类型及其与给定背景数据的区分程度。一些信号始终能与标准模型事件区分开来，并在极低的假阳性率（FPR）下实现较高的ROC曲线下面积（ROC-AUC）和非零的真阳性率（TPR），而另一些信号在低假阳性率下则无法区分。这些信号的差异可归因于不同事件产生的最终状态以及它们与标准模型事件产生的最终状态的接近程度。


我们还注意到，基于ΔR值添加边的GraphSAGE实现方案显著优于其他模型，表明该方法具有广阔的应用前景。在聚合前后引入可学习的变换，可能使模型能够突出某些特征，从而识别尾部出现的异常情况[7]。


此外，我们观察到，有时较高的 ROC-AUC 值在量化模型性能方面可能具有误导性。对于高能物理异常检测而言，FPR 在 1e-5 范围内的情况最符合欧洲核子研究中心 (CERN) 的实时滤波和数据存储限制，但这似乎并未反映在整体 ROC-AUC 得分中。


最后，不同的边构造方法，例如链式、ΔR、能量相关性和k近邻（kNN），似乎对模型的性能有不同的影响；尽管我们注意到，使用ΔR显著提高了GraphSAGE的性能，而能量相关性也提高了GCN模型在所有新信号数据集上的真阳性率（TPR）。能量相关性边可能会放大罕见的相关性，这提高了模型对异常值的敏感性（优于基线方法，即链式），但会降低平均可分离性（如ROC-AUC得分所示）。


未来工作


我们目前的方法仅重构节点特征，探索将目标扩展到重构边结构，可能有助于改进潜在空间编码，从而提升异常检测能力。此外，探索更多基于似然的目标模型，例如变分图编码器或基于能量的评分模型，可以进一步提升模型性能。



参考文献

  • [1]: Zhou, J., et al. (2021). “Graph Neural Networks: A Review of Methods and Applications.” AI Open. https://arxiv.org/abs/1812.08434

  • [2]: Kasieczka, G., et al. (2021). “The LHC Olympics 2020 a community challenge for anomaly detection in high energy physics.” Reports on Progress in Physics, 84(12), 124201. https://arxiv.org/abs/2101.08320

  • [3]: CERN. (2017). “CERN Data Centre passes the 200-petabyte milestone.” https://home.cern/news/news/computing/cern-data-centre-passes-200-petabyte-milestone

  • [4]: ATLAS & CMS Collaborations. (2012). “Observation of a new particle in the search for the Standard Model Higgs boson.” CERN. https://home.cern/news/news/experiments/atlas-and-cms-publish-observations-new-particle

  • [5]: Kipf, T. N., & Welling, M. (2017). “Semi-supervised classification with graph convolutional networks.” ICLR. https://arxiv.org/abs/1609.02907

  • [6]: Fey, M., & Lenssen, J. E. (2019). “Fast graph representation learning with PyTorch Geometric.” ICLR Workshop. https://github.com/pyg-team/pytorch_geometric

  • [7]: Hamilton, W., Ying, Z., & Leskovec, J. (2017). “Inductive representation learning on large graphs.” NeurIPS. https://arxiv.org/abs/1706.02216

  • [8]: Xu, K., et al. (2019). “How powerful are graph neural networks?” ICLR. https://arxiv.org/abs/1810.00826

  • [9]: Veličković, P., et al. (2018). “Graph attention networks.” ICLR. https://arxiv.org/abs/1710.10903

  • [10]: Aarrestad, T., et al. (2022). “The dark machines anomaly score challenge: benchmark data and model independent event classification for the large hadron collider.” SciPost Physics, 12(1), 043. https://arxiv.org/abs/2105.14027

微信群

内容中包含的图片若涉及版权问题，请及时与我们联系删除