Paper Reading #8: DIFUSCO

本文将精读论文 “DIFUSCO: Graph-based Diffusion Solvers for Combinatorial Optimization”，作者 Sun et al.，时间 2023 年，链接 arXiv:2302.08224。

论文概述

这篇论文发表在 NeurIPS 2023 上，提出了 Graph-based Diffusion Solvers for Combinatorial Optimization（DIFUSCO）这一模型，创新性地将扩散模型 (Diffusion Model) 这一在图像和文本生成领域大放异彩的技术，引入到了解决组合优化问题中，并且取得了惊人的效果，尤其是在大规模问题上刷新了 SOTA。

与传统的从零开始一步步构造解不同，DIFUSCO 从纯粹的噪声中，通过学习逐步去噪的过程，完成生成高质量解的任务。这是一个全新的范式。

Abstract

Neural network-based Combinatorial Optimization (CO) methods have shown promising results in solving various NP-complete (NPC) problems without relying on hand-crafted domain knowledge. This paper broadens the current scope of neural solvers for NPC problems by introducing a new graph-based diffusion framework, namely DIFUSCO. Our framework casts NPC problems as discrete $\{0, 1\}$-vector optimization problems and leverages graph-based denoising diffusion models to generate high-quality solutions. We investigate two types of diffusion models with Gaussian and Bernoulli noise, respectively, and devise an effective inference schedule to enhance the solution quality. We evaluate our methods on two well-studied NPC combinatorial optimization problems: Traveling Salesman Problem (TSP) and Maximal Independent Set (MIS). Experimental results show that DIFUSCO strongly outperforms the previous state-of-the-art neural solvers, improving the performance gap between ground-truth and neural solvers from 1.76% to 0.46% on TSP-500, from 2.46% to 1.17% on TSP-1000, and from 3.19% to 2.58% on TSP-10000. For the MIS problem, DIFUSCO outperforms the previous state-of-the-art neural solver on the challenging SATLIB benchmark. Our code is available at https://github.com/Edward-Sun/DIFUSCO.

文章提出了一种全新的 ML4CO 范式。它将 CO 问题映射到 $\{0, 1\}$ 向量，并使用基于图的扩散模型进行去噪，得到高质量的解。模型在 TSP 和 MIS 问题上进行了测试，其中 TSP 问题的规模达到了 10000。

Introduction

作者首先介绍了现有 ML4CO 的三大流派及其局限性：

• 自回归构造式方法 (Autoregressive)：问题是扩展性差。它们需要一步步生成解，每一步都需要进行一次复杂的神经网络计算，导致在处理大规模问题时速度非常慢。
• 非自回归构造式方法 (Non-autoregressive)：问题是表达能力受限。它们通常假设决策变量之间条件独立，一次性生成 heatmap。虽然速度快，但可能难以捕捉解中复杂的多模态结构。
• 改进式启发方法 (Improvement Heuristics)：问题是训练不稳定且扩展困难。它们通常基于 RL，在巨大的解空间中搜索，奖励信号稀疏，导致训练困难且耗时。

为此，作者受到 Diffusion Model 生成概率分布的启发，提出了 DIFUSCO，在图上应用迭代去噪声的过程。作者先将 CO 问题规范为了生成 $\{0, 1\}$ 向量，使用 GNN 做 Encoder，从充满噪声的解中进行去噪。DIFUSCO 旨在同时克服以上所有缺点：既能并行去噪（解决扩展性问题），又能通过迭代优化来捕捉复杂解结构（解决表达能力问题），并且采用稳定高效的训练方式（解决 RL 训练难题）。

模型及训练方法

问题定义

作者将一个组合优化问题的解 $\mathbf{x}$ 统一表示为一个 $\{0, 1\}$ 二进制向量。

• 对于 TSP: $\mathbf{x}$ 是一个长度为边数的向量，$x_i = 1$ 表示第 $i$ 条边被选中
• 对于 MIS: $\mathbf{x}$ 是一个长度为顶点数的向量，$x_i = 1$ 表示第 i 个节点被选中

采用监督学习，损失函数定义为：

$$L(\boldsymbol{\theta})=\mathbb{E}_{s \in \mathcal{S}}\left[-\log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{s *} \mid s\right)\right]$$

这样，求解一个 CO 问题就等价于生成一个满足特定约束且成本最低的 $\{0, 1\}$ 向量 $\mathbf{x}$。

扩散模型：前向过程

前向的加噪过程是固定的、无需学习的。它定义了如何将一个清晰的、真实的解 $\mathbf{x}_0$ 逐步变成纯噪声。

其过程是：我们从一个已知的最优解（或高质量解）$\mathbf{x}_0$ 开始。然后，我们定义 $T$ 个时间步。在每一个时间步 $t$，我们都在 $\mathbf{x}_{t-1}$ 的基础上加入一点点噪声，得到 $\mathbf{x}_t$。经过 $T$ 步之后，最终得到的 $\mathbf{x}_T$ 会变成一个完全随机的向量（纯噪声），与原始解 $\mathbf{x}_0$ 无关。

前向过程可以描述为：

$$q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_0\right)=\prod_{t=1}^T q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right)$$

文章探索了两种加入的噪声类型。

离散扩散：伯努利噪声

其中第一种是伯努利噪声，其前向过程为：

$$q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right)=\operatorname{Cat}\left(\mathbf{x}_t ; \mathbf{p}=\mathbf{x}_{t-1} \mathbf{Q}_t\right)$$

其中

$$\mathbf{Q}_t=\left[\begin{array}{cc} \left(1-\beta_t\right) & \beta_t \\ \beta_t & \left(1-\beta_t\right) \end{array}\right]$$

这个过程可以理解为，在每一步，以一定的概率 $\beta_t$ 随机翻转 $\mathbf{x}_{t-1}$ 中的一些比特位。

进而，前 $t$ 步加噪的过程就可以被写作：

$$q\left(\mathbf{x}_t \mid \mathbf{x}_0\right)=\operatorname{Cat}\left(\mathbf{x}_t ; \mathbf{p}=\mathbf{x}_0 \overline{\mathbf{Q}}_t\right)$$

其中 $\overline{\mathbf{Q}}_t=\mathbf{Q}_1 \mathbf{Q}_2 \ldots \mathbf{Q}_t$。

第 $t-1$ 步的后验概率分布由贝叶斯公式给出：

$$\begin{aligned} & q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0\right)=\frac{q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}, \mathbf{x}_0\right) q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_0\right)}{q\left(\mathbf{x}_t \mid \mathbf{x}_0\right)} \\ & =\operatorname{Cat}\left(\mathbf{x}_{t-1} ; \mathbf{p}=\frac{\mathbf{x}_t \mathbf{Q}_t^{\top} \odot \mathbf{x}_0 \overline{\mathbf{Q}}_{t-1}}{\mathbf{x}_0 \overline{\mathbf{Q}}_t \mathbf{x}_t^{\top}}\right) \end{aligned}$$

这个公式会在计算反向过程及损失的时候都会用到。

连续扩散：高斯噪声

第二种噪声是高斯噪声，这是对离散数据使用连续的扩散过程。

首先，需要将 $\{0, 1\}$ 向量 $\mathbf{x}_0$ 映射到 $\{-1, 1\}$。然后在每一步，向向量中加入一点点均值为 0 的高斯噪声。其前向过程可以描述为：

$$q\left(\hat{\mathbf{x}}_t \mid \hat{\mathbf{x}}_{t-1}\right):=\mathcal{N}\left(\hat{\mathbf{x}}_t ; \sqrt{1-\beta_t} \hat{\mathbf{x}}_{t-1}, \beta_t \mathbf{I}\right)$$

类似地，前 $t$ 步加噪的过程就可以被写作：

$$q\left(\hat{\mathbf{x}}_t \mid \hat{\mathbf{x}}_0\right):=\mathcal{N}\left(\hat{\mathbf{x}}_t ; \sqrt{\bar{\alpha}_t} \hat{\mathbf{x}}_0,\left(1-\bar{\alpha}_t\right) \mathbf{I}\right)$$

其中 $\alpha_t=1-\beta_t$ 并且 $\bar{\alpha}_t=\prod_{\tau=1}^t \alpha_\tau$。这是我们在标准的扩散模型中比较熟悉的公式。

类似地，第 $t-1$ 步的后验概率分布由贝叶斯公式给出：

$$q\left(\hat{\mathbf{x}}_{t-1} \mid \hat{\mathbf{x}}_t, \mathbf{x}_0\right)=\frac{q\left(\hat{\mathbf{x}}_t \mid \hat{\mathbf{x}}_{t-1}, \hat{\mathbf{x}}_0\right) q\left(\hat{\mathbf{x}}_{t-1} \mid \hat{\mathbf{x}}_0\right)}{q\left(\hat{\mathbf{x}}_t \mid \hat{\mathbf{x}}_0\right)}$$

注意，最终还需要将预测的 $\hat{\mathbf{x}}_0$ 转换到 $\{0, 1\}$ 向量。

扩散模型：反向过程

反向的去噪过程是需要学习的，也是整个模型的核心。它的目标是逆转上述的加噪过程。具体来说，需要训练一个神经网络（论文中使用 GNN），让它在只看到更嘈杂的 $\mathbf{x}_t$ 时，学会预测 $\mathbf{x}_{t-1}$ 的样子，即学习 $p_{\boldsymbol{\theta}}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t, s\right)$。

反向过程可以描述为：

$$p_{\boldsymbol{\theta}}\left(\mathbf{x}_{0: T}\right)=p\left(\mathbf{x}_T\right) \prod_{t=1}^T p_{\boldsymbol{\theta}}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)$$

但在两种具体的噪声下，网络要直接预测的目标有所不同。

伯努利噪声下的预测

在离散的扩散模型下，神经网络会预测 $p_{\boldsymbol{\theta}}\left(\widetilde{\mathbf{x}}_0 \mid \mathbf{x}_t\right)$，并由此计算 $p_{\boldsymbol{\theta}}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t, s\right)$：

$$p_{\boldsymbol{\theta}}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)=\sum_{\widetilde{\mathbf{x}}} q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \widetilde{\mathbf{x}}_0\right) p_{\boldsymbol{\theta}}\left(\widetilde{\mathbf{x}}_0 \mid \mathbf{x}_t\right)$$

即训练 GNN 直接从嘈杂的 $\mathbf{x}_t$ 中预测出最原始、最干净的解 $\mathbf{x}_0$。

高斯噪声下的预测

在连续的扩散模型下，神经网络预测每一步加入的高斯噪音：

$$\widetilde{\boldsymbol{\epsilon}}_t=\left(\hat{\mathbf{x}}_t-\sqrt{\bar{\alpha}_t} \hat{\mathbf{x}}_0\right) / \sqrt{1-\bar{\alpha}_t}=f_{\boldsymbol{\theta}}\left(\hat{\mathbf{x}}_t, t\right)$$

再由一个 $\hat{\mathbf{x}}_0$ 的点估计计算 $p_{\boldsymbol{\theta}}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t, s\right)$：

$$p_{\boldsymbol{\theta}}\left(\hat{\mathbf{x}}_{t-1} \mid \hat{\mathbf{x}}_t\right)=q\left(\hat{\mathbf{x}}_{t-1} \mid \hat{\mathbf{x}}_t, \frac{\hat{\mathbf{x}}_t-\sqrt{1-\bar{\alpha}_t} f_{\boldsymbol{\theta}}\left(\hat{\mathbf{x}}_t, t\right)}{\sqrt{\bar{\alpha}}_t}\right)$$

模型训练与推理

损失函数

采用扩散模型中经典的变分下界作为损失函数：

$$\begin{aligned} \mathbb{E}\left[-\log p_{\boldsymbol{\theta}}\left(\mathbf{x}_0\right)\right] \leq \mathbb{E}_q & {\left[-\log \frac{p_{\boldsymbol{\theta}}\left(\mathbf{x}_{0: T}\right)}{q_{\boldsymbol{\theta}}\left(\mathbf{x}_{1: T} \mid \mathbf{x}_0\right)}\right] } \\ =\mathbb{E}_q\left[\sum_{t>1} D_{K L}\right. & {\left[q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0\right) \| p_{\boldsymbol{\theta}}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)\right] } \\ & \left.-\log p_{\boldsymbol{\theta}}\left(\mathbf{x}_0 \mid \mathbf{x}_1\right)\right]+C \end{aligned}$$

可以这么理解这个损失：

• $D_{KL}[ q || p_θ ]$: 最小化这个 KL 散度，就是在让模型的去噪步骤 $p_{\boldsymbol{\theta}}$ 去模仿和逼近那个理想的、有标准答案的去噪步骤 $q$。
• - \log p_{\boldsymbol{\theta}}(\mathbf{x}_0 | \mathbf{x}_1): 它是一个重建项 ，要求模型在去噪的最后一步，能够准确地生成原始的、干净的数据 $\mathbf{x}_0$。

基本训练过程

基本推理过程

快速推理

按照扩散模型的标准定义，推理过程需要进行 $T$ 次（例如 1000 次）迭代去噪，每一步都需要进行一次完整的 GNN 前向传播。这对于追求速度的组合优化求解来说，实在是太慢了。

因此，我们可以减少去噪的步数。我们不再使用 $[1, 2, \dots, T]$ 这一完整的步数序列，而是从中选择一个长度为 $M$（$M \ll T$，如 $M = 50$）的递增子序列 $\tau$，其中 $\tau_1 = 1, \tau_M = T$。这样，我们只需要在这 $M$ 个时间步上进行去噪声。

这种技术在扩散模型领域被称为 DDIM (Denoising Diffusion Implicit Models)，它能在大幅减少计算步数的同时，保持较高的生成质量。

当然，如何选择这个子序列 $\tau$ 也是一个关键问题，这被称为调度。论文中比较了两种调度策略：

• 线性 (Linear) 调度: 子序列中的时间步是均匀分布的。例如，每隔 20 步取一个点。
• 余弦 (Cosine) 调度: 子序列中的时间步在开始阶段（$t$ 接近 $T$）更密集，在后期阶段（$t$ 接近 0）更稀疏。

经过实验发现，余弦调度的效果更好。作者的解释是，在去噪的早期阶段，解中含有大量的噪声，此时模型的任务非常困难，需要更精细、更密集的迭代来确定解的宏观结构。一旦宏观结构确定，后期的微调就可以用更大的步子进行。

网络选择

对于用于噪声预测的神经网络，文章选择了带有边门控机制的各向异性 GNN (Anisotropic GNN with edge gating)。这是一个基于消息传递的 GNN，通过多层迭代，聚合邻居节点和边的信息，来更新每个节点和边的特征表示。

解码方法

经过 $M$ 步去噪后，模型输出的 $\mathbf{x}_0$ 是一个概率热力图 (heatmap)，即一个向量，其中每个元素的值都在 0 到 1 之间，表示对应变量被选中的概率。这个 heatmap 本身并不保证满足 CO 问题的可行性约束。因此，我们需要设计专门的解码策略，将它解码成一个高质量的、离散的、可行的解。

热力图的生成方法是：

• 对于离散扩散，直接使用模型预测的 $p_{\boldsymbol{\theta}}\left(\mathbf{x}_{\mathbf{0}}=1 \mid s\right)$ 作为热力图分数。
• 对于连续扩散，将输出的 $[-1, 1]$ 范围的预测值线性变换到 $[0, 1]$ 范围。

对于 TSP 问题的解码，给定输出的热力图 $\{A_{ij}\}$，有两种解码策略：

• 贪心解码 (Greedy Decoding): 对所有边按 $\left(A_{i j}+A_{j i}\right) /\left\|\mathbf{c}_i-\mathbf{c}_j\right\|$ 排序，从分数最高的边开始，依次尝试将边加入到解中。如果加入的这条边不会和已有的边形成环，则加入；否则，跳过。完成后，可使用 2-opt 等局部搜索算法进行优化。
• 蒙特卡洛树搜索 (MCTS): 将热力图作为先验知识，来指导 MCTS 的搜索方向。MCTS 会优先探索那些在热力图上分数高的区域。

对于 MIS 问题，采用贪心解码。此外，还可以采用多重采样机制，运行多次推理过程选择成本最低的解。

实验及其结果

作者在 TSP 和 MIS 问题上进行了实验。TSP 采用 uniform 分布，节点数为 50 到 10000，使用 Concorde 和 LKH-3 求解器标注。

MIS 采用 SATLIB 和 ER-[700-800]，使用 KaMIS 求解器标注。

TSP 和 MIS 问题上表现都很出色。作者还进行了泛化能力的测试，DIFUSCO 在不同规模的 TSP 问题上表现出良好的泛化能力。