论文笔记DIGRESS: DISCRETE DENOISING DIFFUSION FORGRAPH GENERATION

S20H24

已于 2024-03-25 20:10:16 修改

阅读量3.1k

点赞数 40

文章标签：论文阅读

于 2023-12-17 22:34:13 首次发布

本文链接：https://blog.csdn.net/ZeroXback/article/details/135045357

版权

代码：

GitHub - cvignac/DiGress: code for the paper "DiGress: Discrete Denoising diffusion for graph generation"

一些扩散模型的文章：

An In-Depth Guide to Denoising Diffusion Probabilistic Models – From Theory to Implementation

扩散模型1：基本原理 - 知乎

扩散模型 (Diffusion Model) 简要介绍与源码分析 - 知乎

目录

Abstract

1 Introduction

2 Diffusion models

3 DISCRETE DENOISING DIFFUSION FOR GRAPH GENERATION

3.1迭代过程与逆去噪迭代

3.2去噪网络参数化

3.3等效特性

4利用MARGINAL PROBABILITIES、STRUCTURAL FEATURES改进DiGress

4.1 CHOICE OF THE NOISE MODEL

4.2 STRUCTURAL FEATURES AUGMENTATION

5 CONDITIONAL GENERATION条件生成

EXPERIMENTS

GENERAL GRAPH GENERATION

代码

Pytorch Lightning库

论文代码

PredefinedNoiseScheduleDiscrete类

apply_noise

sample_p_zs_given_zt

compute_batched_over0_posterior_distribution后验分布计算

forward

Abstract

DiGress是用于生成具有分类节点和边属性的离散去噪扩散模型，通过增删边、改变类别逐步编辑有噪声的图形。训练图Transformer网络恢复此过程，将图的分布学习问题简化为节点和边分类任务。本文通过引入马尔科夫噪声模型（在扩散时保留边缘分布的节点和边类型以及结合辅助图论特征）进一步提高样本质量。本文还提出一个训练图级特征生成的程序。

1 Introduction

去噪扩散模型对扩散轨迹降噪，通过对噪声采样和递归降噪产生新样本。扩散模型在图像和视频上优于其他方法。其成功应用提高了构建图生成模型的可能，但是由于其无序性和稀疏性，生成图具有挑战性。连续扩散会使去噪网络难以捕获数据的结构属性。本文对DiGress提出了几个算法上的增强，包括使用在扩散时保留边缘分布的节点和边类型的噪声模型，介绍了一种新的基于图级属性的图生成指导方法，用辅助结构和光谱特征来增强我们的去噪网络的输入。从噪声图导出的这些特征使用噪声模型的离散性质，与基于高斯的模型相比保留了稀疏的噪声图。这些改进增强了DiGress在各种图形生成任务上的性能。

2 Diffusion models

去噪扩散模型包括两个主要部分：噪声模型、去噪神经网络。

噪声模型q逐渐破环数据点x以创建一系列噪声数据点 $(z^1,...,z^T)$ 。它有马尔科夫结构: $q(z^1,...,z^T|x)=q(z^1|x)\prod{_{t=2}^T}q(z^t|z^{t-1})$ 。

去噪网络从 $z^t$ 预测 $z^{t-1}$ 进行反转。噪声从先验分布中被采样，然后通过去噪网络的迭代应用反转。

扩散模型需满足：

1.分布 $q(z^t|x)$ 有closed-form公式以允许在不同时间步长上并行训练。

2.后验 $p_\theta(z^{t-1}|z^t)=\int q(z^{t-1}|z^t,x)dp_\theta(x)$ 应有closed-form表达式，这样x可以作为神经网络的目标。

3.极限分布 $q_\infty=lim_{T\rightarrow\infty}q(z^T|x)$ 不依赖x，以将其作为先验分布推理

当使用高斯噪声可满足以上性质。当需要对分类数据建模时，可通过将数据嵌入到具有类别的独热编码的连续空间中使用高斯噪声。但是高斯噪声是图的差的噪声模型，因为破坏稀疏性以及连通性等图论概念。因此离散扩散更适合图生成任务。

噪声由转换矩阵 $(Q^1,...,Q^T)$ 表示， $[Q^t]_{ij}$ 表示从状态i跳到状态j的概率： $q(z^t|zt-1)=Z^t-1Q^t$

去噪网络从 $G^t$ 学习干净图。预测的分布用 $q(G^{t-1}|G,G^t)$ 组合以计算 $p_\theta(G^t-1|G^t)$ 以及采样一个离散的 $G^{t-1}$ 。

3 DISCRETE DENOISING DIFFUSION FOR GRAPH GENERATION

模型处理具有分类节点和边属性的图，分别由空间X和E表示，基数为a和b。

3.1迭代过程与逆去噪迭代

模型分别对每个节点和边特征进行扩散。转移概率： $[Q{_X^t}]_{ij}=q(x^t=j|x^{t-1}=i)$ , $[Q{_E^t}]_{ij}=q(e^t=j|e^{t-1}=i)$ 。添加噪声产生 $G^t=(X^t,E^t)$ 即从分类分布中对每个节点和边类型进行采样。

DiGress第二个部分是由 $\theta$ 参数化的去噪神经网络。它以噪声图 $G^t=(X^t,E^t)$ 为输入预测干净图，通过每个节点和边的预测概率及真实图优化交叉熵损失。使用网络预测估计反向扩散迭代 $p_\theta(G^{t-1}|G^t)$ 采样新图。建模分布为节点和边上的乘积：

$p_\theta(G^{t-1}|G^t)=\prod \limits_{1\leq i\leq n}p_\theta(x{_i^{t-1}}|G^t)\prod \limits_{1\leq i,j\leq n}p_\theta(e{_ij^{t-1}}|G^t)$

这个分布用于对离散 $G^{t-1}$ 进行采样， $G^{t-1}$ 将在下一个时间步作为去噪网络的输入。

3.2去噪网络参数化

去噪网络以噪声图 $G^t=(X^t,E^t)$ 为输入，输出表示干净图上的预测分布张量 $X^` E^`$ 。使用Transformer网络操作图级特征y，注意力机制是边预测的自然模型。它首先使用自关注更新节点特征，使用FiLM层合并边特征和全局特征,使用未归一化的注意力分数更新边特征，使用池化的节点和边特征更新图级特征。时间步长归一化为[0，1]，并将其视为y内部的全局特征。

3.3等效特性

图对于节点的重新排序是不变的，这意味着n！矩阵可以表示相同的图。为了有效地从这些数据中学习，设计出不需要用随机排列来扩充数据的方法是至关重要的。这意味着如果训练数据被置换，梯度更新不应该改变。->需要：a permutation equivariant architecture and a permutation invariant loss。DiGress满足这两个属性。

Lemma 3.1 置换同变性：输出顺序与输入顺序对应。（置换不变性：输出与输入顺序无关）

Lemma 3.2 不变损失：损失函数可分解为 $\sum _il_X(p{_i^X},x_i)+\sum _{i,j}l_E(p{_i^E},e_{ij})$ ，在每个节点和每条边上分别计算的两个函数 $l_X$ 和 $l_E$ 是置换不变的。

Lemma 3.3可交换性：DiGress产生可交换的分布，即它生成具有节点特征X和邻接矩阵A的图，对于任何置换π，邻接矩阵A满足 $P(X,A)=P(\pi^TX,\pi^TA\pi)$ 。

4利用MARGINAL PROBABILITIES、STRUCTURAL FEATURES改进DiGress

4.1 CHOICE OF THE NOISE MODEL

图是稀疏的因此边类型的边缘分布远远不是均匀。为了改进均匀转换，假设使用接近真实数据分布的先验分布使训练更容易。先验分布不能任意选择，因为它需要是置换不变的，以满足交换律（引理3.3）。因此该分布的自然模型是所有节点的单个分布u和所有边的单个分布v的乘积。

4.2 STRUCTURAL FEATURES AUGMENTATION

DiGress在离散空间上操作，其噪声图是不完整的，允许在每个扩散步骤计算各种图描述符。这些描述符可以输入到网络中，以帮助去噪过程，从而产生用于训练DiGress并从中采样的算法1和2。包含这些额外的特征可以通过实验提高性能，但不是必需的。应考虑选择包含哪些特征及其计算的计算性能。

5 CONDITIONAL GENERATION条件生成

文章（Equivariant diffusion for molecule generation in 3d）执行条件生成是使用目标属性来训练去噪网络，但当条件属性改变时需要重新训练模型。本文使用一个被训练用于预测一个干净图的目标属性的regressor $g_\eta$ 。这个回归量通过在每个采样步骤调整预测的分布并将其推向具有所需属性的图来指导无条件扩散模型。

EXPERIMENTS

GENERAL GRAPH GENERATION

评估正确建模图的各种属性的能力，如所生成的图形是否是统计上可区分的SBM模型，或者如果它们是平面和连接。

代码

Pytorch Lightning库

包含五个部分：

Computations(init) 初始化相关计算
Train Loop(training_step) 每个step中执行的代码，相当于训练过程中处理一个batch的内容
Validation Loop(validation_step) 在一个epoch训练完以后执行Valid
Test Loop(test_step) 在整个训练完成以后执行Test
Optimizer(configure_optimizers) 配置优化器等

测试部分中每个batch的训练被称为一个step，故先运行train_step函数。当经过多个batch进行验证，运行validation_step函数。完成一个epoch的训练以后，会对整个epoch结果进行验证，运行validation_epoch_end函数。如果需要的话，可以调用测试部分代码。在training_step 、validation_step、test_step 后面都紧跟着其相应的training_step_end(self，batch_parts)和training_epoch_end(self, training_step_outputs) 函数，当然，对于校验和测试，都有相应的*_step_end和*_epoch_end函数。

论文代码

metrics：离散版digress用的是一种交叉熵代价函数

models：一个graph transformer，当成一个输入X、E、y和一个节点mask，输出更好的X、E、y的模型

PredefinedNoiseScheduleDiscrete类

调用diffusion_utils.cosine_beta_schedule_discrete或custom_beta_schedule_discrete方法生成对应的噪声数组betas。

diffusion_utils.cosine_beta_schedule_discrete:计算离散的余弦调度噪声系数

forward方法：t_normalized表示归一化的时间步，t_int表示整数时间步，相乘计算得到的整数时间步，从betas中获取对应的噪声系数并返回。

apply_noise

随机生成时间步，使用PredefinedNoiseScheduleDiscrete类的alpha_bar计算状态转移概率【Qt = prod(1 - beta_t) * I + (1 - prod(1 - beta_t)) * K】probX、probE。对probX、probE进行离散特征采样，获取采样后的特征值。然后将特征值转换为one-hot编码。

sample_p_zs_given_zt

从条件分布 $p(z_s,z_t|z_0)$ 采样，计算得到当前时间步和目标时间步的噪声水平，获取转换模型中的转移矩阵，根据给定的输入数据和转移矩阵，计算条件概率分布 $p(z_s,z_t|z_0)$ 的后验分布。对预测结果和后验分布进行加权处理，得到加权后的预测结果。对加权后的预测结果进行归一化处理，得到概率分布。使用采样函数从概率分布中采样出隐状态，将采样得到的隐状态转换为one-hot编码形式。根据给定的节点掩码对输出结果进行掩码操作。