🔥 神经扩散过程(NDP)

【摘 要】 高斯过程为指定函数的先验分布和后验分布提供了一个优雅的框架。然而，它们计算代价很高，并且受到协方差函数表达能力的限制。我们提出了神经扩散过程 (NDP)，这是一种基于扩散模型的新方法，它学习如何从函数的分布中进行采样。使用一种新颖的注意力块，我们能够将随机过程的性质（例如可交换性）直接整合到神经扩散架构中。经验表明，神经扩散过程能够捕获与高斯过程的真实贝叶斯后验接近的函数型分布。这使各种下游任务成为可能，包括超参数的边缘化和贝叶斯优化。

【原 文】 Dutordoir, V. et al. (2022) ‘Neural Diffusion Processes’. arXiv. Available at: http://arxiv.org/abs/2206.03992 (Accessed: 21 February 2023).

1 简介

高斯过程 (GP) 提供了一个强大的框架来定义函数的分布 ^[23]。高斯过程的一个关键优势是其可交换性，这使得从训练模型中获得任意位置的一致预测成为可能。但是，有两个关键因素限制了高斯过程在现实世界中可处理的任务范围。首先，推断成本过高，为 $\mathcal{O}(N^3)$，其中 $N$ 是数据点的数量。其次，传统仅对超参数进行点估计的方法通常存在问题，但是获得完整的边缘分布在计算上又过于昂贵 ^[13]。

基于神经网络的生成式模型也能够学习函数的分布，从而使其能够模仿高斯过程。成功的方法包括神经过程 (NP)的元学习 ^[8] 、高斯神经过程 ^[2][18] 和基于 VAE 的 ^[19][6] 方法。通过在训练期间利用大量小型数据集，这些方法能够在预测时贯穿数据集传输知识。使用神经网络很有吸引力，因为大部分计算工作都花费在训练过程中，而预测任务通常变得更加直接。基于神经网络的方法的另一个主要优点是允许对高斯过程的超参数进行完整边缘化（即便对于相对较小的数据集），这对于传统高斯过程模型来说在计算上具有挑战性。

本文通过扩展现有基于神经网络的生成模型来改进上述方法。“概率去噪扩散模型” ^[27][28][9] 最近被证明在图像 ^[20][22]、分子结构 ^[34]、点云 ^[17] 和音频信号^[12] 的样本生成效果明显。但现成的扩散模型无法从连续的随机过程中生成样本。这非常难，最大的难点在于：（1）随机过程的样本是一个函数，一个可以在输入域的任意位置进行计算的函数；（2）对于从随机过程中抽取的样本，往往被要求遵守某些性质（如可交换性：当函数计算的顺序改变时，联合概率分布不会改变）。这使得该问题从根本上不同于图像或音频生成。

本文贡献：

• 我们提出了一种新模型，并称之为 神经扩散过程 (NDP)，它将扩散模型的应用扩展到了随机过程，并且能够描述函数的分布。通过对模型扩散的随机变量进行索引，神经扩散过程将扩散模型推广到了无限维的函数空间。
• 我们特别注意地将随机过程的对称性和其他性质（包括可交换性）强制纳入模型中，以改进训练过程。这些性质是在一种被称为 “二维注意力块” 的新型注意力组件帮助下实施的，该方法能够保证 输入维度 和 序列（即数据点）排序 上的 等方差性。
• 我们通过将神经扩散过程应用于一系列贝叶斯推断任务（包括先验采样和条件采样、回归、超参数的边缘化、贝叶斯优化）来展示神经扩散过程的能力和多功能性。
• 我们还提出了一种使用神经扩散过程的新型全局优化方法。

2 随机过程的去噪扩散概率模型

扩散模型 ^[9][27][28] 能够从复杂的高维分布中创建高质量样本。到目前为止，文献主要集中在从一维、二维和三维（如音频、图像、点云和分子）的分布中生成样本。此外，在所有这些应用中生成的特征，要么是一个在预定义输入上的可变化输出，要么是在一个可变化输入位置上的固定输出。例如，在图像生成应用中，我们作出了像素的预定义配置假设。在音频合成中，随机变量根据声线的离散化进行排序。

相对来说，神经扩散过程（NDP）将扩散模型推广到了无限维函数空间，其方法是允许用另一组随机变量 $\boldsymbol{x}$ 来索引函数观测值 $\boldsymbol{y}$。这种方法使得在输入域中的任何位置处，预测神经扩散过程的函数样本成为可能。该方法同样适用于对随机过程（例如高斯过程）进行建模。

扩散模型有 正向过程 和 逆向过程 两个运行组件，如 图 1 所示。 正向过程由一个马尔可夫链组成，它逐渐向数据中添加随机噪声。逆向过程的任务是反转此链，以便从噪声中重建所需的数据样本。为了准确地近似复杂的逆向过程，在 第 3 节 中，我们将为神经扩散过程介绍一种新颖的神经网络架构，并讨论它如何将随机过程的已知对称性和性质（如可交换性）强制连接到架构中。

图 1：NDP 正向过程和逆向过程的示意图。 （从左到右）在正向过程中，从函数（蓝线）中抽取的数据点（蓝点）逐层受到增加的噪声扰动，直到在 $T$ 层数据点（即包括输入也包括输出）变成了白噪声。（从右到左）逆向过程包含逐渐去除噪声的各层，直到我们到达一个数据点。

首先，我们将扩展传统扩散模型组件以适应对函数的采样。

（1）符号表示

令 $[a]_i$ 为一个列向量的第 $i$ 个元素。设 $f : \mathcal{X} \rightarrow \mathbb{R}$ 是一个函数，$\boldsymbol{y} \in \mathbb{R}^n$ 是 $N$ 个含噪的函数值向量，$\boldsymbol{x} \in \mathcal{X}^N$ 是包含 $N$ 个输入的向量，两者之间遵循关系 $[\boldsymbol{y}]_i = f([\boldsymbol{x}]_i) + η$ , 其中 $η \sim \mathcal{N}(0, \sigma^2)$ 对应于观测噪声。

（2）正向过程

对于神经深度过程（NDP），数据分布 $(\boldsymbol{x_0, y_0}) \sim q(\boldsymbol{x_0, y_0})$ 由随机过程的样本组成。我们可以通过定义两个先验来构造它，一个是输入位置的先验 $q(\boldsymbol{x}_0)$ ，另一个是条件观测的先验 $q(\boldsymbol{y_0 | x_0})$。来自这些先验的样本将被用作训练数据集。

我们以独立同分布方式在输入域中均匀地抽取输入；而观测的先验则采用中心化的高斯过程，即 $q(\boldsymbol{y_0 | x_0}) = \mathcal{N}(\boldsymbol{y_0}; \mathbf{0}, k_ψ(\boldsymbol{x_0, x_0}) + \sigma^2 \mathbf{I})$，其中协方差函数 $k_ψ$ 的超参数为 $ψ$。

然后我们将正向过程定义为一条固定的马尔可夫链，它通过在每一步中添加少量噪声来迭代地破坏样本数据 $(\boldsymbol{x_0, y_0})$

$$q(\boldsymbol{y}_{0:T}, \boldsymbol{x}_{0:T}) = q(\boldsymbol{x}_0) q(\boldsymbol{y}_0 | \boldsymbol{x}_0) \prod^{T}_{t=1} q ( \begin{bmatrix} \boldsymbol{x}_t \\ \boldsymbol{y}_t \end{bmatrix} \mid \begin{bmatrix} \boldsymbol{x}_{t-1} \\ \boldsymbol{y}_{t-1} \end{bmatrix} ,t ) \tag{1}$$

其中破坏分布 $q(· |·, t)$ 是联合高斯分布:

$$q ( \begin{bmatrix} \boldsymbol{x}_t \\ \boldsymbol{y}_t \end{bmatrix} \mid \begin{bmatrix} \boldsymbol{x}_{t-1} \\ \boldsymbol{y}_{t-1} \end{bmatrix} ,t) = \mathcal{N}( \begin{bmatrix} \boldsymbol{x}_t \\ \boldsymbol{y}_t \end{bmatrix} \mid \sqrt{1 − β_t} \begin{bmatrix} \boldsymbol{x}_{t-1} \\ \boldsymbol{y}_{t-1} \end{bmatrix} , β_t \mathbf{I} ) \tag{2}$$

上述每个步骤中的噪声大小由预先指定的方差列表 $\{β_t \in (0, 1)\}^T_{t=1}$ 控制。

注意，正向过程中不涉及任何学习，它只是创建了一系列随机变量 $\{(\boldsymbol{x_t, y_t})\}^{T}_{t=1}$，随着 $t$ 接近 $T$，它逐渐看起来更像白噪声，如 图 1 所示。

正向过程的一个有用特性是：鉴于高斯具有线性叠加的特点，我们可以绕过迭代的正向生成过程，直接在任意时间步 $t$ 处采样，$t \in \{1, 2, \ldots , T\}$。令 $γ_t = \sqrt{1 − \bar{α_t}}$ 和 $\bar{α}_t = \prod^{t}_{j=1}(1 − β_j)$。依靠重参数化技巧，可以使用 $\boldsymbol{x}_t$ 和 $\boldsymbol{y}_t$ 进行采样，即：

$$\boldsymbol{x}_t = \sqrt{\bar{α}_t} \boldsymbol{x}_0 + \gamma_t \boldsymbol{\varepsilon_{x}} \\ \boldsymbol{y}_t = \sqrt{\bar{α}_t}\boldsymbol{y}_0 + \gamma_t \boldsymbol{\varepsilon_{y}} \tag{3}$$

其中 $\boldsymbol{\varepsilon_{x}, \varepsilon_{y}} \sim \mathcal{N}(\mathbf{0, I})$。

（3）逆向过程

逆向过程的目标是从时间 $T$ 的白噪声开始创建目标数据（即函数的抽取）。正向过程虽然被设计为马尔可夫的，但真实的逆向概率需要完整序列。让我们将这个逆向过程的近似定义为：

$$p_θ(\boldsymbol{x}_{0:T},\boldsymbol{y}_{0:T}) = p(\boldsymbol{x}_T , \boldsymbol{y}_T) \prod^{T}_{t=1} p_θ( \begin{bmatrix} \boldsymbol{x}_{t−1}\\ \boldsymbol{y}_{t−1} \end{bmatrix} \mid \begin{bmatrix} \boldsymbol{x}_t\\ \boldsymbol{y}_t \end{bmatrix} ,t )$$

它从白噪声初始状态 $p(\boldsymbol{x}_T , \boldsymbol{y}_T ) = \mathcal{N}(\mathbf{0,I})$ 开始，其中 $p_\theta$ 是参数化的逆向马尔可夫核。逆向核的任务是学习反转 式 (1) 中的每个步骤。这是一个复杂的函数，取决于输入 $\boldsymbol{x}_t、\boldsymbol{y}_t$ 以及时间。借鉴 Ho 等 ^[9]，我们将 $p_\theta$ 参数化为高斯分布，其中均值用当前状态的参数化输出 $\boldsymbol{\epsilon}^{x}_{\theta}(·)$ 和 $\boldsymbol{\epsilon}^{y}_{\theta}(·)$ 构造：

$$p_\theta ( \begin{bmatrix} \boldsymbol{x}_{t−1} \\ \boldsymbol{y}_{t−1} \end{bmatrix} \mid \begin{bmatrix} \boldsymbol{x}_t\\ \boldsymbol{y}_t \end{bmatrix} ,t ) = \mathcal{N}( \begin{bmatrix} \boldsymbol{x}_{t−1} \\ \boldsymbol{y}_{t−1} \end{bmatrix} \mid \frac{1}{\sqrt{1 − β_t}}( \begin{bmatrix} \boldsymbol{x}_t \\ \boldsymbol{y}_t \end{bmatrix} − \frac{β_t}{\gamma_t} \begin{bmatrix} \boldsymbol{\epsilon}^{x}_{\theta}(\boldsymbol{x}_t, \boldsymbol{y}_t, t)\\ \boldsymbol{\epsilon}^{y}_{θ}(\boldsymbol{x}_t, \boldsymbol{y}_t, t) \end{bmatrix}), \frac{γ^2_t}{γ^2_{t−1}} β_t \mathbf{I}) \tag{4}$$

在接下来的内容中，我们将 $\boldsymbol{\epsilon}^{x}_{\theta}(·)$ 和 $\boldsymbol{\epsilon}^{y}_{\theta}(·)$ 称为噪声模型，他们会有效地学习和预测添加到真实数据中的噪声，进而可以在每个逆向步骤中根据模型将其剔除。在 第 3 节 中，我们讨论了噪声模型的设计，一般通过逆向过程产生的函数样本，能够遵循随机过程的某些性质。不过对于本节的其余部分，我们暂时假设噪声模型的参数化形式已经给定，且有一组需要优化的参数 $θ$。

噪声模型 $\boldsymbol{\epsilon}^{x}_{\theta}(·)$ 和 $\boldsymbol{\epsilon}^{x}_{\theta}$ 的参数 $θ$ 可以通过优化模型证据 $\log p_\theta(\boldsymbol{x}_0, \boldsymbol{y}_0)$ 的变分下界来训练。不过，借鉴 Kingma 等 ^[11]，神经扩散过程使用重新加权的变分下界，其中权重选择是为了样本的质量和更容易优化（详细推导见 附录 A），即损失函数为：

$$\mathcal{L}_\theta = \mathbb{E}_{\boldsymbol{x}_0,\boldsymbol{y}_0,\boldsymbol{\varepsilon_{x}},\boldsymbol{\varepsilon_{y}},t} \left [ \| \boldsymbol{\varepsilon_{x}} − \boldsymbol{\epsilon}^{x}_{\theta}(\boldsymbol{x}_t, \boldsymbol{y}_t, t) \|^2 + \| \boldsymbol{\varepsilon_{y}} −\boldsymbol{\epsilon}^{y}_{\theta}(\boldsymbol{x}_t, \boldsymbol{y}_t, t) \|^2 \right] \tag{5}$$

我们可以通过模拟（采样）方法获得对 $\mathcal{L}_\theta$ 的随机估计，即对 $\boldsymbol{x}_0 \sim q(\boldsymbol{x}_0)$、$\boldsymbol{y}_0 \sim q(\boldsymbol{y}_0 | \boldsymbol{x}_0)$、$\boldsymbol{\varepsilon_{x}}$、$\boldsymbol{\varepsilon_{y}} \sim \mathcal{N}(\mathbf{0, I})$ 和 $t \sim \mathcal{U}({1, 2, \ldots , T})$ 进行采样，并将样本用于蒙特卡洛积分。

请注意，$\boldsymbol{x}_t = \sqrt{\bar{α}_t}\boldsymbol{x}_0 + \gamma_t \boldsymbol{\varepsilon_{x}}$ 和 $\boldsymbol{y}_t = \sqrt{\bar{α}t}\boldsymbol{y}_0 + \gamma_t \boldsymbol{\varepsilon_{y}}$ 是使用 式(3) 创建的。

很明显，训练噪声模型 $\boldsymbol{\epsilon}^{x}_{\theta}(·)$、$\boldsymbol{\epsilon}^{y}_{\theta}(·)$ 的目的，是用来预测添加到真实数据点的噪声 $\boldsymbol{\varepsilon_{x}}$、$\boldsymbol{\varepsilon_{y}}$。问题仍然是为 $\boldsymbol{\epsilon}^{x}_{\theta}(·)$ 和 $\boldsymbol{\epsilon}^{y}_{\theta}(·)$ 采用哪种参数化模型；我们同样将此问题的讨论推迟到 第 3 节，并同时讨论如何在准确近似噪声模型的假设下获得相干样本。

图 2：条件样本：蓝色曲线是基于不同概率模型的上下文数据集（黑点）的后验样本。我们还绘制了黑色样本的经验均值和两倍标准差虚线。更多示例在 图 D.1 中给出。

2.1 推断和条件样本

到目前为止，我们已经展示了神经扩散过程如何能够通过使用逆向核 $p_\theta$ 遍历逆向过程来生成先验 $p(\boldsymbol{x}_0, \boldsymbol{y}_0)$ 的函数样本抽取。在本节中，我们将详细介绍神经扩散过程如何能够从后验 $p(\boldsymbol{x}_0, \boldsymbol{y}_0 | \mathcal{D})$ 中抽取条件样本，其中 $\mathcal{D} = \{([\boldsymbol{x}^c_0]_i \in \mathcal{X} , [\boldsymbol{y}^c_0]_i \in \mathbb{R})\}^{n}_{i=1}$ 是上下文数据集。这类似于我们如何从以数据集 $\mathcal{D}$ 为条件的高斯过程中抽取后验样本，如 图 2a 所示。

为了符号简单起见，我们提取了单一样本 $(x_0,y_0)$ 的采样方案，但需要强调的是，所有操作都能够以批处理形式获得 $M$ 个相关输入和来自同一函数抽取的相应输出。

根据扩散模型 ^[16] 图像修复应用的最新进展，我们调整逆向过程以考虑 $\mathcal{D}$ 中的信息（注：高斯场中某个位置的空间预测，可以被视为一种修复）。我们希望条件样本与上下文数据集相关并保持一致。条件采样过程如下：

我们首先增强中间状态以将上下文数据集 $(\boldsymbol{x}_t^c, \boldsymbol{y}_t^c)$ 中的信息和 $(\boldsymbol{x}_t, \boldsymbol{y}_t)$ 中的初始随机状态信息包含在内。这会产生增强后的中间状态 $\boldsymbol{x}_t = [\boldsymbol{x}_t^c, \boldsymbol{x}_t]$ 和 $\boldsymbol{y}_t = [\boldsymbol{y}_t^c, \boldsymbol{y}_t]$。其中上下文数据集信息 $(\boldsymbol{x}_t^c, \boldsymbol{y}_t^c)$ 对应的中间状态组分，来自于从上下文数据 $(\boldsymbol{x}^c_0, \boldsymbol{y}^c_0)$ 开始的正向过程（式 (6.a)），而 $(x_t, y_t$ 组分则创建于从初始状态 $t = T$、$\boldsymbol{x}_T \sim \mathcal{N}(\mathbf{0, I})$ 、 $\boldsymbol{y}_T \sim \mathcal{N}(0, 1)$ 开始的逆向核（式（6.b））：

$$\begin{bmatrix} \boldsymbol{x}_t^c\\ \boldsymbol{y}_t^c \end{bmatrix} \sim \mathcal{N}( \begin{bmatrix} \boldsymbol{x}_t^c\\ \boldsymbol{y}_t^c \end{bmatrix} \mid \sqrt{\bar{α}t} \begin{bmatrix} \boldsymbol{x}^c_0\\ \boldsymbol{y}^c_0 \end{bmatrix}, γ^2_t \mathbf{I}_n ) \qquad \text{第 t 个正向步骤 (式 (3))} \tag{6.a}$$

$$\begin{bmatrix} x_t \\ y_t \end{bmatrix} \sim p_\theta ( \begin{bmatrix} x_t\\ y_t \end{bmatrix} \mid \begin{bmatrix} \boldsymbol{x}_{t+1} \\ \boldsymbol{y}_{t+1} \end{bmatrix},t+1) \qquad \text{第 t 个逆向步骤 (式 (4))} \tag{6.b}$$

该方案确保：对于每个 式 (6.b) 的逆向步骤，我们能够利用包含在增强状态 $(\boldsymbol{x}_t, \boldsymbol{y}_t)$ 中的上下文数据集信息。

我们在 图 2c 和 图 5 中给出了此过程的一维和二维例子，并在 附录 B 中给出了伪代码。

请注意，抽取条件样本只需要一个预训练的无条件神经深度过程。在训练期间，神经扩散过程只关心学习先验知识。不需要为了创建训练样本做不方便的数据集拆分等工作，类似于神经过程（ ^[8][7][10]）。通过设计，相同的神经扩散过程将从零泛化到任意数量的条件点，如 图 2c 所示。

2.2 确定性的输入位置

在监督学习任务中，由于输入是预先知道的，所以习惯上条件分布 $p(\boldsymbol{y}_0 | \boldsymbol{x}_0)$ 比联合分布 $p(\boldsymbol{x}_0, \boldsymbol{y}_0)$ 更受重视。如果这是模型的真实使用场景，则我们建议不要对输入 $\boldsymbol{x}_0$ 做扩散，而是保持固定的输入，即 $\boldsymbol{x}_0 = \boldsymbol{x}_1 = \ldots = \boldsymbol{x}_T$。从数学上讲，这等同于确定性地将 $\boldsymbol{\varepsilon_{x}}$ 设置为零，从而使其相应的预测 $\boldsymbol{\epsilon}^{x}_{\theta}(·)$ 变得不必要，进而将 式(5) 简化为：

$$\mathcal{L}_\theta = \mathbb{E}_{\boldsymbol{x}_0,\boldsymbol{y}_0,\boldsymbol{\varepsilon_{y}},t} \left[ \|\boldsymbol{\varepsilon_{y}} −\boldsymbol{\epsilon}^{y}_{\theta}(\boldsymbol{x}_0, \boldsymbol{y}_t, t)\|^2 \right] \tag{7}$$

噪声模型 $\boldsymbol{\epsilon}^{y}_{\theta}(·)$ 现在作用于未损坏的输入 $\boldsymbol{x}_0$ 和损坏的观测 $\boldsymbol{y}_t = \sqrt{\bar{α}t}\boldsymbol{y}_0 + \gamma_t \boldsymbol{\varepsilon_{y}}$。如前所述，对于经典的回归任务，扩散输入位置会带来不必要的复杂性。

不过，在 第 4 节 中，我们将介绍一种新颖的贝叶斯优化算法，该算法中的一个示例需要访问完整分布 $p(\boldsymbol{x}_0, \boldsymbol{y}_0)$，而且至关重要。

最后，在训练中使用 式 (7) 仍然会产生一个具有无限维函数空间先验的模型，它可以在域 $\mathcal{X}$ 的任何地方进行计算。

3 噪声模型和二维注意力块

在本节中，我们回顾了神经扩散过程的噪声模型架构和关键组件，包括第 3.2 节中的一种新型注意力模块。如上一节所述，噪声模型产生的输出是 $\boldsymbol{\epsilon}_x(·)$ (相对于 $\boldsymbol{\epsilon}^{y}_{\theta}(\cdot)$），它试图预测需要被减去的噪声 $\boldsymbol{\varepsilon_{x}}$ ( 相对于 $\boldsymbol{\varepsilon_{y}}$ ) 以从 $\boldsymbol{x}_t$ 中恢复去除噪声后的值 $\boldsymbol{x}_0$（以及从 $\boldsymbol{y}_t$ 中恢复 $\boldsymbol{y}_0$）。

神经扩散过程使用神经网络作为噪声模型。原则上，可以使用任何神经网络。但如果我们希望神经扩散过程模仿高斯过程，则噪声模型必须学会生成函数的先验分布。

一般来说，我们期望这样的先验能够拥有几个关键的对称性，而这些对称性将严重影响我们对体系结构的选择。

• 可交换性：正如 Kolmogorov 展开定理所指出的：随机过程的最基本性质之一是不同输入顺序的等方差性。因此，打乱数据集 $\mathcal{D}$ 中数据点的顺序或待预测的顺序，不应当影响数据的概率（即数据是可交换的）。
• 我们还期望输入维度的排序具备不变性。虽然后验分布不需要在输入维度上具有对称性，但先验分布不包含手头特定问题的信息，因此不能偏向于任何一个输入。虽然可以通过在训练期间操纵损失函数或通过数据增强来鼓励神经网络显示不变性，但更可取的方法是将其硬连接到体系结构中，从一开始就强制执行不变性。

用于噪声模型的神经网络架构如 图 3 所示。

图 3：噪声预测模型的架构，可以用于神经扩散过程中的每个步骤。灰色框表示二维注意力块，如 第 3.2 节 所述

3.1 数据集大小和维度不可知性

如 图 3 所示，神经扩散过程的噪声模型在 $B$ 个批量数据集上运行，其中每个数据集由 $N$ 个输入和相应的输出组成。我们假设输入是 $D$ 维的（即 $\mathcal{X} = \mathbb{R}^D$）并且输出是一维的（$\mathcal{Y} = \mathbb{R}^1$ ）。

神经扩散过程架构必须具备的一个有用性质，即该架构与数据集大小 $N$ 和维度 $D$ 无关。也就是说，网络权重不能依赖于 $N$ 或 $D$，并且输出应当简单地匹配相应输入的张量形状。

为了实现此功能，神经扩散过程首先将输入 $(\boldsymbol{x}_t \in \mathbb{R}^{N×D}, \boldsymbol{y}_t \in \mathbb{R}^n )$ 预处理为形状 $[N, D, 2]$（注：后文将省略批次 $B$，因为各批次上的所有操作都相同 ）。此张量形状允许后续的二维注意力块（接下来解释）也是 $N$ 和 $D$ 的。预处理步骤包括将一维输出 $\boldsymbol{y}_t$ 复制 $D$ 次，然后将它们与每个维度的相应输入组合在一起。这给出了 $D$ 个配对 $\{[x^d, y] \in \mathbb{R}^2\}^{D}_{d=1}$，如 图 3 中的红色框所示，附录 E 也列出了相应 Python 代码。预处理步骤之后是一个全连接层，该层将 $2D$ 的输入投影到 $H$ 维空间中。

对于另一个输入（时间 $t$)，我们遵循 Vaswani 等 ^[32] 的方法，使用循环 $128$ 维编码（详见 附录 D.1），然后使用密集层将其投影到相同的 $H$ 维空间。

该模型与数据集大小 $N$ 和维度 $D$ 的无关性要求非常重要，因为其会带来重要结果：不需要为规模或维度不同的数据集分别训练不同的神经深度过程。这使得训练一个能够处理不同 $N$ 或 $D$ 值下游任务的统一模型成为可能，如 第 4 节 所示。

3.2 二维等变注意力块

我们新提出的二维注意力块是神经扩散过程模型的核心。

它在输入 $(\boldsymbol{x}_t, \boldsymbol{y}_t)$ 和 $t$ 经过预处理并嵌入到 $H$ 维空间后起作用。该组件的核心由两个多头自注意力 (MHSA ^[32]) 层组成。 多头自注意力层作用于输入的两个不同轴：其中一个关注输入维度轴，而另一个关注整个数据集序列轴。两者的输出随后相加并通过一个非线性函数，然后通过残差连接（图 3 中的虚线）将输出反馈回下一个二维注意力块。此过程会被重复迭代执行多次，如 图 3 所示。

我们现在证明这个块关于维度和数据集序列的排序具有等变性（即，如果打乱输入的行和/或列，则会以相同方式打乱输出）。令 $\mathbf{A}_t : \mathbb{R}^{N×D×H} \rightarrow \mathbb{R}^{N×D×H}, \mathbf{s} \mapsto \mathbf{A}_t(\mathbf{s})$ 表示二维注意力块。它作用于从 $(\boldsymbol{x}_t, \boldsymbol{y}_t)$ 创建的输入 $\mathbf{s}_t$，如 图 3 所示。以下命题成立：

【命题 1】 设 $\prod_N$ 是索引 $\{1, \ldots，N\}$ 的所有排列构成的集合。令 $\mathbf{s} \in \mathbb{R}^{N×D×H}$ ，并且 $(π_n \circ \mathbf{s}) \in \mathbb{R}^{N×D×H}$ 表示一个张量，其第一维（即 $N$ 对应的维度）中索引的排序由 $π_n \in \prod_N$ 给出。则对于所有排列 $π_n \in \prod_N$，有:

$$π_n \circ \mathbf{A}_t(\mathbf{s}) = \mathbf{A}_t(π_n \circ \mathbf{s})$$

【命题 2】 设 $\prod_D$ 是索引 $\{1, \ldots , D\}$ 的所有排列构成的集合。令 $\mathbf{s} \in \mathbb{R}^{N×D×H}$，并且 $(π_d \circ \mathbf{s}) \in \mathbb{R}^{N×D×H}$ 表示一个张量，其第二维（即 $D$ 对应的维度）中的索引排序由 $π_d \in \prod_D$ 给出。则对于所有排列 $π_d \in \prod_D$：

$$π_d \circ \mathbf{A}_t(\mathbf{s}) = \mathbf{A}_t (π_d \circ \mathbf{s})$$

证明：在 附录 C 中包含了【命题 1】 和【命题 2】的证明。

鉴于排列 $π_n$ 和 $π_d$ 作用在张量的不同轴上，我们在 附录 C 中证明了它们是兼容的且可交换的。这使得二维注意力块 $\mathbf{A}_t$ 同时等变于数据集的顺序及其输入维度。

3.3 噪声模型的输出：维度不变，序列顺序等变

神经扩散过程的噪声模型，将会输出逆向过程的两个关键要素：$\boldsymbol{\epsilon}^{x}_{\theta}(·)$ 将预测添加到输入中的噪声； $\boldsymbol{\epsilon}^{y}_{\theta}(·)$ 则会预测添加到输出中的噪声。两者都是通过对多个双向注意力层的输出求和形成的，如 图 3 架构图中的紫色求和所示。由于此求和是逐元素运算，因此不会破坏上面的等变性质。在这个运算之后，噪声模型进行了分叉，因为我们分别需要 $\boldsymbol{\epsilon}^{x}_{\theta}(·)$ 和 $\boldsymbol{\epsilon}^{y}_{\theta}(·)$ 的不同性质。对于 $\boldsymbol{\epsilon_x}$，我们期望在 $N$ 和 $D$ 中具有等变性；但对于 $\boldsymbol{\epsilon}^{y}_{\theta}$，我们期望在 $D$ 中具有等变性。也就是说，函数 $\boldsymbol{\epsilon}^{y}_{\theta}$ 的输出不应该受到打乱后的数据集特征影响。借鉴 Zaheer 等 ^[35]，我们可以通过对输入维度轴 ($D$) 求和来适应这种不变性，如 图 3 中的绿色求和所示。

总结一下，使用【命题 1】和【命题 2】中的符号，我们在 附录 C 中已经证明，噪声模型遵循：

$$π_n \circ π_d \circ \boldsymbol{\epsilon}^{x}_{\theta}(\boldsymbol{x}_t, \boldsymbol{y}_t, t) = \boldsymbol{\epsilon}^{x}_{\theta}(π_n \circ π_d \circ \boldsymbol{x}_t, π_n \circ \boldsymbol{y}_t, t)$$

以及

$$π_n \circ\boldsymbol{\epsilon}^{y}_{\theta}(\boldsymbol{x}_t, \boldsymbol{y}_t, t) =\boldsymbol{\epsilon}^{y}_{\theta}(π_n \circ π_d \circ \boldsymbol{x}_t, π_n \circ \boldsymbol{y}_t, t)。$$

4 实验评估

在本节中，我们进行实验来说明 神经扩散过程的几个重要特征，包括它们产生高质量条件样本的能力、边缘化核超参数的能力，以及它们的输入维度不变性。然后，我们考虑 神经扩散过程在高维贝叶斯优化 (BO) 和使用 神经扩散过程对联合 $p(\boldsymbol{x}_0, \boldsymbol{y}_0)$ 建模能力的新型全局优化策略中的应用。所有实验共享相同的模型架构，如 图 3 所示。有 $5$ 个二维注意力块。所有模型配置都在附录 D.1 中给出。该代码可以在补充中找到。

数据集所有实验的训练数据都是综合创建的。我们使用高斯过程先验样本，$\boldsymbol{y}_0 | \boldsymbol{x}_0 \sim \mathcal{N}(0, k_ψ(\boldsymbol{x}_0, \boldsymbol{x}_0) + \sigma^2 \mathbf{I})$，具有较小的固定似然方差。高斯过程都在输入域 $[−1, 1]^D$ 上定义。我们专注于固定核 $k_ψ$。长度尺度是固定的或从对数正态分布中采样。数据生成过程的数字细节在 附录 D.1 中给出。

【条件采样】 为了表明神经扩散过程学习了可用于抽取连贯样本的函数的准确分布，我们从一个简单的一维回归任务开始。 图 2 （和补充中的图 D.1）说明了从高斯过程、注意力潜在 NP (ALNP) ^[10] 和我们的 神经扩散过程中抽取的后验样本。 ALNP 和 神经扩散过程在来自长度为 0.2 的平方指数 (SE) 核的先验样本上进行训练。对于 ALNP，我们使用了来自参考实现 ^[3] 的预训练模型。虽然 ALNP 不会遭受传统 NP 中已知的欠拟合问题，但远离数据点的样本行为并不能准确反映其真实函数形式，由高斯过程基本事实给出。来自 神经扩散过程的样本被认为更接近于目标分布。

图 4：核长度尺度边缘化：来自 NDP 的样本，以越来越多的数据点（黑点）为条件，如第一行所示。每个样本的颜色对应于最有可能产生该样本的核的长度尺度。最可能的长度尺度来自具有最高模型证据的相应 GPR 模型。底行说明了所有样本中最有可能确定的长度尺度的分布。随着更多数据点的提供，可能的长度尺度分布会收敛于用于生成数据的长度尺度（橙色虚线）。

【超参数边缘化】 通常，高斯过程模型是通过最大化超参数的可能性来训练的。理想情况下，贝叶斯推断是通过边缘化超参数来执行的，这可以显着提高性能，尽管需要额外的计算成本 ^[13][26]。 神经扩散过程的一个关键功能是它能够在其训练的全部数据范围内生成逼真的条件样本。因此，它实际上可以边缘化它所暴露的不同超参数，甚至不同的核。 图 4 演示了该模型如何从一系列不同的长度尺度生成样本，并且随着观测到更多数据，它会缩小真实长度尺度。该模型使用来自 Matérn- $\frac{3}{2}$ 核的高斯过程样本进行训练，具有从 $\log \mathcal{N}(\log(0.5), \sqrt{0.5})$ 中采样的各种长度尺度。

图 5：2D 回归：来自先验和后验的样本证明了模型的维度不变性。后验样本以观察结果（圆圈）为条件，其中圆圈的颜色对应于模型所依据的地面真实值。

图 6：贝叶斯优化：在基于汤普森抽样的贝叶斯优化中使用的不同概率模型对 3D（左）和 6D（右）问题的遗憾。不确定性估计基于使用不同的随机种子重复实验五次。

【高维 BO】 在本实验中，我们强调神经扩散过程可用于 $D ≥ 3$ 的问题。我们使用离散 Thompson 采样 BO ^[24][5][30] 解决附录 D.3 中给出的 3D 和 6D 最小化问题。在这种情况下，概率模型被用作昂贵的黑盒目标的替代品。在每次迭代中，代理在输入域中的 $128$ 个随机位置进行评估，并选择对应于最小值的输入作为下一个查询点。在此位置评估目标并将其添加到上下文数据集中。 图 6 显示了不同模型的遗憾（与真实最小值的距离）。我们观测到 神经扩散过程几乎与高斯过程回归的性能相匹配，高斯过程回归是此类任务的黄金标准模型。高斯过程回归和 神经扩散过程之间的重要区别在于 神经扩散过程在 BO 循环期间不需要训练，而高斯过程回归在每一步都需要重新训练。

【使用密度估计进行优化】 我们可以利用 神经扩散过程对完整联合分布 $p(\boldsymbol{x}^*, \boldsymbol{y}^* | \mathcal{D})$ 进行建模的能力，构想出一种不同的优化策略，而不是从 $p(y^* | x^*, \mathcal{D})$ 中获取样本).考虑根据我们当前对最小值 $y^*$ 的信念来调整神经扩散过程模型。这使我们能够从 $p(x^* | y^*)$ 中获取样本，其中包含根据模型输入域中最小值所在位置的信息。在 图 7 我们使用这个想法说明了一个全局优化例程。在每一步中，我们从 $p(y^*)$ 中采样一个目标，它描述了我们对最小值的信念。对于该图，我们使用了 $p(y^*)$ 的截断法线，其中范围是使用已经观测到的数据指定的（我们参考 附录 D.4 了解更多详细信息）。通过采样 $p(\boldsymbol{x}^* | \boldsymbol{ y }^*, \mathcal{D})$ 选择下一个查询点，从而系统地寻找全局最小值。这个实验展示了 神经扩散过程模拟功能输入和输出之间复杂交互的能力——这是一项没有明确训练的任务。

图 7：通过扩散输入位置进行全局优化。我们在每个步骤的目标 $y^*$ 值（虚线）上调整模型，并使用 NDP 从 $p(y^* | x^*, \mathcal{D})$ 中采样，即满足此新目标值的输入位置的后验分布.上排的面板说明了优化的进展。在下一行中，蓝色分布对应于模型对目标所在位置的信念，而红色三角形标记选定的查询点。

5 相关工作

我们在引言中简要讨论了相关工作。在这里，我们更深入地研究了相关方法，这些方法侧重于使用神经网络来描述函数的概率分布。

【神经过程族】 神经过程 ^[7][8] 是与 神经扩散过程密切相关的一系列模型，因为它们利用深度神经网络来定义函数上丰富的概率分布。神经过程由两部分组成：编码器和解码器。编码器用于将上下文数据集嵌入到潜在表示中，然后解码器使用它在目标点进行预测。条件神经过程 ^[7] 仅预测后验的边缘，这使得无法从模型中抽取连贯的样本。 Latent 神经过程 ^[8] 试图通过合并潜在变量来缓解这个问题，但这引入了源于近似推断的困难。

尽管神经过程与神经扩散过程具有相似的性质，但它们不包含维度不变性，并且它们通常会遭受对上下文集的欠拟合，导致高估的不确定性以及未通过上下文点的样本。已经提出了对该方法的一些改进，包括基于注意力的神经过程 ^[10] 和 Gaussian 神经过程 ^[2][18]；然而，据我们所知，没有一个被证明可以在 $D > 3$ 中进行推断。

用于加速高斯过程训练和模型选择的神经网络 越来越多的工作证明了深度神经网络架构如何提供一种有前途的方法来分摊与高斯过程相关的计算密集型计算。例如，Liu 等 ^[14] 使用 transformer 来预测谱混合核的超参数，而 Simpson 等 ^[25] 使用 transformers 为给定数据集推荐合适的高斯过程核。

6 结论

我们提出了神经扩散过程 (NDP)，这是一种去噪扩散模型方法，用于学习函数空间的概率，并生成函数的先验和条件样本。神经扩散过程通过允许对随机变量进行索引，将扩散模型推广到无限维函数空间。我们引入了一个新的神经网络构建组件，即二维注意力块，它将维度和序列等变地连接到神经网络架构中，使其表现得像一个随机过程。我们根据经验表明，该模型能够捕获接近真实贝叶斯后验的函数分布。

【局限性和未来工作】与其他扩散模型一样，我们发现神经扩散过程的样本质量随着扩散步骤 $T$ 的数量而提高。然而，相对于 GAN 等其他架构，这确实会导致推断时间变慢。可以结合加速推断过程的技术来改善这个问题。 Watson 等提出的方法。 ^[33] 对于神经扩散过程特别感兴趣，因为我们有一个原则性和可区分的指标来评估样本质量，即相应的高斯过程的边缘似然。最后，参数化傅立叶域中的样本，例如使用 ^[21][4]，可能是一种有趣的替代方法。

参考文献

• [1] J. Berkeley, H. B. Moss, A. Artemev, S. Pascual-Diaz, U. Granta, H. Stojic, I. Couckuyt, J. Qing, N. Loka, A. Paleyes, S. W. Ober, and V. Picheny. Trieste, 2022. URL: https://gith ub.com/secondmind-labs/trieste.
• [2] W. P. Bruinsma, J. Requeima, A. Y. Foong, J. Gordon, and R. E. Turner. The Gaussian neural process. In Advances in Approximate Bayesian Inference, 2021.
• [3] Y. Dubois, J. Gordon, and A. Y. Foong. Neural process family. http://yanndubs.github.io/Neural-Process-Family/.
• [4] V. Dutordoir, N. Durrande, and J. Hensman. Sparse Gaussian processes with spherical harmonic features. In International Conference on Machine Learning, 2020.
• [5] V. Dutordoir, N. Knudde, J. van der Herten, I. Couckuyt, and T. Dhaene. Deep Gaussian Process metamodeling of sequentially sampled non-stationary response surfaces. In Winter Simulation Conference, 2017.
• [6] V. Fortuin, D. Baranchuk, G. Rätsch, and S. Mandt. GP-VAE: Deep probabilistic time series imputation. In International Conference on Artificial Intelligence and Statistics, 2020.
• [7] M. Garnelo, D. Rosenbaum, C. Maddison, T. Ramalho, D. Sa\boldsymbol{x}_ton, M. Shanahan, Y. W. Teh, D. J. Rezende, and S. M. A. Eslami. Conditional Neural Processes. In International Conference on Machine Learning, 2018.
• [8] M. Garnelo, J. Schwarz, D. Rosenbaum, F. Viola, D. J. Rezende, S. M. A. Eslami, and Y. W. Teh. Neural processes. arXiv preprint arXiv:1807.01622, 2018.
• [9] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. In Neural Information Processing Systems, 2020.
• [10] H. Kim, A. Mnih, J. Schwarz, M. Garnelo, S. M. A. Eslami, D. Rosenbaum, O. Vinyals, and Y. W. Teh. Attentive neural processes. In International Conference on Learning Representations, 2019.
• [11] D. P. Kingma, T. Salimans, B. Poole, and J. Ho. Variational diffusion models. In Neural Information Processing Systems, 2021.
• [12] Z. Kong, W. Ping, J. Huang, K. Zhao, and B. Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations, 2020.
• [13] V. Lalchand and C. E. Rasmussen. Approximate inference for fully Bayesian Gaussian process regression. In Advances in Approximate Bayesian Inference, 2020.
• [14] S. Liu, X. Sun, P. J. Ramadge, and R. P. Adams. Task-agnostic amortized inference of Gaussian process hyperparameters. In Neural Information Processing Systems, 2020.
• [15] J. R. Lloyd, P. Orbanz, Z. Ghahramani, and D. M. Roy. Random function priors for exchangeable arrays with applications to graphs and relational data. In Neural Information Processing Systems, 2012.
• [16] A. Lugmayr, M. Danelljan, A. Romero, F. Yu, R. Timofte, and L. Van Gool. RePaint: Inpainting using denoising diffusion probabilistic models. arXiv preprint arXiv:2201.09865, 2022.
• [17] S. Luo and W. Hu. Diffusion probabilistic models for 3d point cloud generation. In Computer Vision and Pattern Recognition, 2021.
• [18] S. Markou, J. Requeima, W. Bruinsma, and R. Turner. Efficient Gaussian Neural Processes for Regression. arXiv preprint arXiv:2108.09676, 2021.
• [19] S. Mishra, S. Flaxman, T. Berah, M. Pakkanen, H. Zhu, and S. Bhatt. piVAE: Encoding stochastic process priors with variational autoencoders. arXiv preprint arXiv:2002.06873, 2020.
• [20] A. Nichol, P. Dhariwal, A. Ramesh, P. Shyam, P. Mishkin, B. McGrew, I. Sutskever, and M. Chen. Glide: Towards photorealistic image generation and editing with te\boldsymbol{x}_t-guided diffusion models. arXiv preprint arXiv:2112.10741, 2021.
• [21] A. Rahimi and B. Recht. Random features for large-scale kernel machines. In Neural Information Processing Systems, 2007.
• [22] A. Ramesh, P. Dhariwal, A. Nichol, C. Chu, and M. Chen. Hierarchical te\boldsymbol{x}_t-conditional image generation with CLIP latents. arXiv preprint arXiv:2204.06125, 2022.
• [23] C. E. Rasmussen and C. K. Williams. Gaussian processes for machine learning. MIT Press, 2006.
• [24] B. Shahriari, K. Swersky, Z. Wang, R. P. Adams, and N. De Freitas. Taking the human out of the loop: A review of Bayesian optimization. Proceedings of the IEEE, 2015.
• [25] F. Simpson, I. Davies, V. Lalchand, A. Vullo, N. Durrande, and C. E. Rasmussen. Kernel identification through transformers. In Neural Information Processing Systems, 2021.
• [26] F. Simpson, V. Lalchand, and C. E. Rasmussen. Marginalised Gaussian processes with nested sampling. In Neural Information Processing Systems, 2021.
• [27] J. Sohl-Dickstein, E. A. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. arXiv preprint arXiv:1503.03585, 2015.
• [28] Y. Song and S. Ermon. Improved techniques for training score-based generative models. In Neural Information Processing Systems, volume 33, 2020.
• [29] UCI Machine Learning Repository. https://archive.ics.uci.edu/ml/index.php.
• [30] S. Vakili, H. Moss, A. Artemev, V. Dutordoir, and V. Picheny. Scalable Thompson Sampling using Sparse Gaussian Process Models. In Neural Information Processing Systems, 2021.
• [31] M. van der Wilk, V. Dutordoir, S. John, A. Artemev, V. Adam, and J. Hensman. A framework for interdomain and multioutput Gaussian processes. arxiv preprint arXiv:2003.01115, 2020.
• [32] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. In Neural Information Processing Systems, 2017.
• [33] D. Watson, W. Chan, J. Ho, and M. Norouzi. Learning fast samplers for diffusion models by differentiating through sample quality. In International Conference on Learning Representations, 2021.
• [34] M. Xu, L. Yu, Y. Song, C. Shi, S. Ermon, and J. Tang. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022.
• [35] M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R. Salakhutdinov, and A. J. Smola. Deep Sets. In Neural Information Processing Systems, 2017.

附录

B 算法

在本节中，我们列出了用于训练和采样 神经扩散过程的伪代码。处理确定性输入位置特殊情况的代码（第 2.2 节）用 “OR” 标记。在这种情况下，输入 $\boldsymbol{x}_t$ 没有中间状态，$y$ 的噪声模型输入只是 $\boldsymbol{x}_0$。

B.1 训练

训练过程涉及拟合 神经扩散过程噪声模型神经网络的参数，输出为 $\boldsymbol{\epsilon}^{x}_{\theta}$，$\boldsymbol{\epsilon}^{x}_{\theta}$，以便它们可以准确地预测出添加到受损函数样本中的噪声。在每次迭代中，我们创建一批数据集，每个数据集都来自高斯过程的抽取。我们强调数据点的数量甚至数据集的输入维数在训练过程中可能会有所不同。

B.2 先验采样和条件抽样

实际上，对先验进行采样的代码是 “当上下文数据集为空时” 条件采样代码的特例。但为了说明的清楚，我们在这里将两者都列出来。

条件采样算法基于利用扩散模型进行图像修复的一些最新成果 ^[16] 。在 算法 3 中，为了获得最佳结果，我们遵循 Lugmayr 等 ^[16] 的方法，并使用将预测与上下文数据集混合的重采样步骤。在所有的实验中，$U$ 都被设置为 $5$。

E 代码

代码可在 https://github.com/vdutor/neural-diffusion-processes 处获得。

预处理方法

神经扩散过程的噪声模型，首先将输入 $\boldsymbol{x}_t \in \mathbb{R}^{N×D}$ 和 $\boldsymbol{y}_t \in \mathbb{R}^N$ 重新排列为 $[N, D, 2]$ 形状的张量，如第 3.1 节所述。下面我们展示了此预处理的代码：

import tensorflow as tf 

def preprocess(x: tf.Tensor , y: tf.Tensor) -> tf.Tensor:
    """ 
    Transform inputs to split out the x dimensions for dimension agnostic processing. 

    :param x: [B, N, D] 
    :param y: [B, N, 1] 
    :return: [B, N, D, 2] 
    """ 

    D = tf.shape(x)[-1] 
    x = tf.expand_dims(x, axis=-1) 
    y = tf.repeat(tf.expand_dims(y, axis=-1), D, axis=2) 
    return tf.concat([x, y], axis=-1)