Stanford CS224W: Machine Learning with Graphs
By Yiwen Chen, Aleksandr Timashov, and Yue (Andy) Zhang as part of the Stanford CS224W course project.
动机
生成模型（generative models）（尤其是基于分数的生成模型）的最新进展已在各个领域取得了重大进展。然而，图生成模型（graph generation models）的一个关键挑战仍然存在：缺乏置换不变性（permutation invariance）。这篇博文介绍了 EDP-GNN 架构，这是一种用于图生成的置换不变方法。EDP-GNN 一旦在图数据集上进行训练，就可以使用退火朗之万动态采样（Langevin dynamics sampling）生成反映输入分布的新图。
在基于分数（score-based）的生成模型中，我们专注于学习对数数据密度（log data density）的梯度（分数函数）。这种方法允许更具表现力的架构，因为它绕过了对分数规范化的需要。通过利用图神经网络等固有的置换不变结构，我们确保等效邻接矩阵（equivalent adjacency matrices）的输出一致。EDP-GNN 架构就是一个例子，它在保持置换不变性的同时学习图分布分数。我们的讨论还将涵盖基于分数的图生成中训练和采样的实际方面。
图神经网络（Graph Neural Networks）
使用案例
图神经网络 (GNN) 是机器学习领域的一项重大创新，专门用于分析网络或图等结构化数据（data structured）。这使得它们成为各种应用的理想选择。这包括社交网络和蛋白质折叠（proteins fold）方式等，它们自然地构成了互连点（interconnected points）。GNN 在交通管理（traffic management）等领域至关重要，它们通过了解复杂的道路网络来帮助优化路线。它们还在金融领域发挥着重要作用，通过分析交易网络（transaction networks）来检测欺诈行为（fraud detection），并在推荐系统中根据用户的兴趣网络推荐产品。GNN 的突出之处在于它们能够解释这些数据集中错综复杂的关系网，而传统神经网络则认为这是一项具有挑战性的任务。通过利用数据中的连接和模式的力量，GNN 为这些不同领域提供了富有洞察力和有效的解决方案，展示了它们在处理基于图的信息方面的多功能性和重要性。
置换不变性（Permutation invariance）
GNN 与其他神经网络架构的一个主要区别是它们能够保持置换不变性。这意味着即使输入图中节点的顺序发生变化，GNN 的输出也不会改变。此属性对于图数据至关重要，因为图中节点的顺序通常是任意的，并且不会传达有意义的信息。在具有 N 个节点的图中，最多可以有 N！（阶乘factorial）个不同的邻接矩阵表示同一张图。节点置换的不变性使 GNN 能够更好地泛化并应用于各种图结构而不会降低有效性。GNN 在处理多样化和复杂的图结构数据方面的多功能性和适应性凸显了它们在人工智能领域日益增长的重要性。
基于分数的生成模型
基于分数的图分析模型使用神经网络来建模一个称为分数函数（score function）的独特方面。该函数本质上表示图数据分布对数的梯度。使用分数函数的优势在于其简单性：与需要复杂归一化以确保概率总和为 1 的直接密度建模不同，分数函数绕过了这一要求。
这些模型围绕一系列噪声（noise）levels构建，表示为σ ₁、σ ₂、…、σ ₗ，每个level逐渐变小（σ ₁ < σ ₂ < … < σ ₗ）。该模型以这些噪声level为条件，使用它们逐渐塑造输出。具体来说，该模型从噪声开始并应用一种称为退火朗之万动力学（annealed Langevin dynamics）的技术来生成邻接矩阵。这个过程是理解基于分数的模型如何有效处理图数据的关键。
朗之万动态采样（Langevin dynamic sampling）
一旦在图数据集上训练了 EDP-GNN 模型，我们就可以使用退火的朗之万动力学从模型中采样新图。
对于给定的 EDP-GNN 得分模型s(x, σ)，用于采样的朗之万更新为：
从高斯N(0, I)中采样噪声zₜ
更新xₜ₊₁ = xₜ + α/2 s(x, σ) + √(α) zₜ
我们看到更新过程中有两个关键部分。项α/2 s(x, σ)是梯度更新，将xₜ移向学习数据分布中更可能的区域。项√(α) zₜ为采样添加了一些噪声——如果没有它，样本将收敛到局部最优，将无法获得正确的样本。此过程运行大量迭代T，而步长 alpha 应该足够小。
退火的朗之万动力学
为了将朗之万采样程序扩展到退火朗之万采样，我们利用了 EDP-GNN 得分模型的噪声调节。我们首先使用普通的朗之万程序从非常嘈杂的分布s(x, σ_L)中进行采样。然后使用生成的样本初始化从s(x, σ_{L-1}) 进行的采样，这提供了稍微少一些噪声的梯度更新。我们逐渐降低噪音并细化样本，直到s(x, σ₁ )具有输入数据分布的最准确表示。
退火朗之万动态采样的完整伪代码如下：
Initialize x_1 from a simple prior distribution, eg. Gaussian
For sigma_i in sigma_L, …, sigma_1:
Define the step size alpha_i = epsilon * sigma_i^2 / sigma_L^2
For t in 1, … T:
Sample z_t from N(0, I)
Update x_{t+1} = x_t + alpha/2 * s(x, sigma) + sqrt(alpha) z_t
Initialize for the next noise level x_1 = x_t
Return x_T在 pytorch_geometric 中，EDP-GNN（和其他基于分数的图生成模型）的退火朗之万动态采样在 torch_geometric.utils.langevin 中实现。你可以像这样使用它：
import torch_geometric.utils.langevin as langevin
score_model = … # EDP-GNN model
node_features = … # num_nodes x feature_dim matrix of node features
adjs, node_flags = langevin.generate_initial_sample(
batch_size=5, num_nodes=10
)
def score_func(adjs, node_flags):
return score_model(node_features, adjs, node_flags)
new_adjs = langevin.sample(
score_func, adjs, node_flags, num_steps=5000, quantize=True
)在 pytorch_geometric 实现中，对退火的朗之万动力学过程进行了一些针对图的修改：
通过此过程生成的邻接矩阵具有连续值，但在实践中，我们经常希望对具有离散边值的图进行采样。我们可以在最后做到这一点，只需定义一个阈值来二值化或量化邻接矩阵值即可。
由于邻接矩阵是对称的，初始邻接矩阵 x_1 和噪声项 z_t 也必须对称。
噪声调度器（Noise Scheduler）
噪声调度器是一种常用实用程序，用于在评分模型和其他基于扩散的模型中生成噪声levels。在基于评分的生成建模框架中，在训练过程中，系统地将噪声以各种levels添加到数据中，表示为 sigma 值。这种方法对于学习噪声条件评分模型至关重要。我们已经实现了噪声调度器，如“通过估计数据分布的梯度进行生成建模”中所述，它专门以对数尺度生成噪声levels。
def get_smld_sigma_schedule(
sigma_min: float,
sigma_max: float,
num_scales: int,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
) -> Tensor:
r"""Generates a set of noise values on a logarithmic scale for "Score
Matching with Langevin Dynamics" from the `"Generative Modeling by
Estimating Gradients of the Data Distribution"
<https://arxiv.org/abs/1907.05600>`_ paper.
This function returns a vector of sigma values that define the schedule of
noise levels used during Score Matching with Langevin Dynamics.
The sigma values are determined on a logarithmic scale from
:obj:`sigma_max` to :obj:`sigma_min`, inclusive.
Args:
sigma_min (float): The minimum value of sigma, corresponding to the
lowest noise level.
sigma_max (float): The maximum value of sigma, corresponding to the
highest noise level.
num_scales (int): The number of sigma values to generate, defining the
granularity of the noise schedule.
dtype (torch.dtype, optional): The output data type.
(default: :obj:`None`)
device (torch.device, optional): The output device.
(default: :obj:`None`)
"""
return torch.linspace(
math.log(sigma_max),
math.log(sigma_min),
num_scales,
dtype=dtype,
device=device,
).exp()要使用此功能，必须指定调度所需的 sigma 值的范围和数量。这些参数的确定取决于模型和数据的具体需求。例如，Niu 2020 中使用 EDP-GNN 进行小规模图生成使用六个噪声levels，而图像生成模型可能需要数百或数千个噪声levels。它通常涉及一个实验和调整过程，以确定给定数据集和模型架构的最有效调度。
扩散模型中使用的噪声调度器与 Langevin Dynamics 的Score Matching 中的噪声调度程序在概念上相似，如论文“Denoising Diffusion Probabilistic Models”中所述。但是，它们的用途不同，并且采用不同的机制。在扩散模型中，调度器会生成一系列 beta 值，用于前向扩散过程。虽然我们的项目没有直接使用这种特定类型的噪声调度器，但我们已经实现了一些变体，以促进未来基于扩散的图模型。
边密集（Edgewise Dense）预测图神经网络（EDP-GNN）
多通道 GNN 层
多通道 GNN 层是图同构网络 (Graph Isomorphism Network，GIN) 层的扩展，它采用了一种新颖的消息传递方法。其核心思想是同时在相同的节点特征上运行消息传递算法，但使用不同的邻接矩阵。此过程最终产生的输出是所有 C 个通道的输出的串联，通道数等于中间邻接矩阵的数量。此设置中第 m 层的公式针对 C 通道邻接矩阵输入量身定制，如下所示：
值得注意的是，多通道 GNN 层是最终 EDP-GNN 模型的基本组成部分，使其能够高效处理复杂的图结构。
实现的主要部分是消息和forward函数。
def message(self, x_j, edge_weight):
return x_j * edge_weight
def forward(self, x: Union[Tensor, OptPairTensor], edge_indices: List[Adj],
edge_weights: List[Tensor], size: Size = None) -> Tensor:
...
edge_weights_cat = torch.cat(edge_weights, dim=-1)[:, None]
# duplicate features to run over C channels
...
# propagate_type: (x: OptPairTensor)
out_cat = self.propagate(edge_index_cat, edge_weight=edge_weights_cat,
x=x_cat, size=size)
x_r = x_cat[1]
if x_r is not None:
out_cat = out_cat + (1 + self.eps) * x_r # N * C x F_in
# reshape to get concatenated features for each of C channels
out_cat = out_cat.reshape((C, N, -1)).permute((1, 0, 2))
out = out_cat.reshape((N, -1))
return self.nn(out)EDP-GNN层
EDP-GNN 层的功能是处理由 C 通道邻接矩阵和节点特征组成的输入，并输出更新的 C' 通道邻接矩阵（也包含节点特征）。
使用多通道 GNN 层进行节点特征的计算。随后，生成一个新的 C' 通道邻接矩阵。此过程涉及将新计算的节点特征与原始 C 通道邻接矩阵结合使用。该方法使用连接和多层感知器 (MLP) 网络集成这些元素。这种方法强调了该层在转换和更新图结构方面的效率。
在这种对无向图进行建模的方法中，保持预测邻接矩阵的对称性至关重要。为此，我们将预测矩阵与其转置版本相加。此过程确保最终的邻接矩阵准确反映图的无向性质。为了更全面地理解这种转换，下面提供了详细的公式。
最终的网络架构
最终的网络由三层表示，如我们在图像中看到的一样。在输入到我们的 EDP-GNN 模型之前，图要经过预处理。这涉及使用双通道邻接矩阵：一个通道用于原始邻接矩阵，另一个通道用于其取反对应项，所有条目均反转。此外，我们根据节点的加权度初始化节点特征。此过程可确保图以最佳格式供 EDP-GNN 模型处理。
具有 3 层的高级 EDP-GNN
具体来说，如果我们从数据中得到节点特征X，那么每个节点的初始值如下。
最后要注意的一点是，我们网络中输入和输出的维度是相同的。这一点至关重要，因为网络旨在对评分函数进行建模。保持一致的维度可确保网络准确地表示和处理评分函数，这对于模型的有效性至关重要。
噪声level调节
如前所述，有必要根据不同的噪声level对网络的每一层进行调节。我们通过引入一些可学习的参数（具体来说，是通过增加增益和偏差（gains and biases））来实现这种噪声调节。条件多层感知器 (MLP) 层表示为：
其中 α 和 β 是每个噪声level σ 的可学习参数。
结论
在这篇文章中，我们探索了基于分数的图生成方法的训练和采样过程。该方法利用图神经网络 (GNN) 来保留生成的图中置换不变性的关键归纳偏差。此外，我们演示了如何使用 PyG 实用程序实现 EDP-GNN 网络。如果您想深入了解技术细节，我们鼓励您阅读https://arxiv.org/abs/2003.00638上的原始论文。
参考
[1] C. Niu, Y. Song, J. Song, S. Zhao, A. Grover, and S. Ermon, Permutation Invariant Graph Generation via Score-Based Generative Modeling, (2020), AISTATS 2020
[2] Song, Y. and Ermon, S. (2019). Generative modeling by estimating gradients of the data distribution. arXiv preprint arXiv:1907.05600
[3] Xu, K., Hu, W., Leskovec, J., and Jegelka, S. (2018a). How powerful are graph neural networks? arXiv preprint arXiv:1810.00826
[4] Jonathan Ho, Ajay Jain, Pieter Abbeel. Diffusion Probabilistic Models https://arxiv.org/abs/2006.11239
[5] Yang Song. Diffusion and Score-Based Generative Models YouTube video. https://youtu.be/wMmqCMwuM2Q?si=64WID5_82zywM6_C
[6] https://github.com/ermongroup/GraphScoreMatching/tree/master
