🔥 神经过程作为函数的分布

【原文】 Kaspar Martins, 2018, Neural Processes as distributions over functions, https://kasparmartens.rbind.io/post/np/

在 ICML2018 中，提出了一些关于神经过程的有趣工作。见论文《条件神经过程》和同一作者在研讨会上提出的《神经过程》的后续工作。

神经过程（NPs）引起了我的注意，因为它们本质上是一种基于神经网络（NN）的概率模型，可以表示随机过程的分布。所以神经过程结合了两个世界的元素:

深度学习：神经网络是灵活的非线性函数，可以直接训练。
高斯过程：高斯过程提供了一个概率框架，用于学习一类广泛的非线性函数的分布。

两者各有其优缺点。在数据有限的体系中，高斯过程由于其概率性质和捕获不确定性的能力更受欢迎。这与（非贝叶斯）神经网络不同，后者表示一个单一函数，而不是一个关于函数的分布。但是，在存在大量数据的情况下，后者可能更可取，因为训练神经网络在计算上比推断高斯过程更容易实现可扩展性。

神经过程旨在结合这两个世界的优点。

神经过程背后的想法很有趣，但我觉得自己缺乏直觉，也对神经过程作为函数的先验缺乏足够深的理解。但我相信，理解一个东西的最好方法往往是实现它，在简单问题上进行经验性的尝试，并最终向别人解释这一点。因此，本文是我对神经过程的回顾、讨论和尝试。

在阅读我的文章之前，建议读者先看一下原始论文。我在这里重点讨论了《神经过程》，但你可能会发现从《条件神经过程》可能会更容易一些，因为它本质上是神经过程的非概率版本。

1 什么是神经过程？

神经过程是一种基于神经网络表示函数上分布的方法。下图说明了神经过程模型建立以及训练的大致思路。

Fig01

图 1: 神经过程的顶层视图

给定一组观测值 $\{(x_i, y_i)\}$ ，将其分为两个集合：“背景集 $\{(x_c,y_c)\}, c=1,\ldots,C$ ” 和 “目标集 $\{(x_t,y_t)\}, t=1,\ldots ,T$ ”。给定 背景集 中的已观测数据 $(x_c, y_c)$ 和 目标集 中未观测过的输入 $x_t^{\ast}$ ，我们的目标是预测出相应的函数值 $y_t^{\ast}$ 。

我们可以把神经过程想成是： 以 背景集 为条件对 目标集 建模。信息通过潜在空间 $z$ ，从 背景集 （图 1 左边）流向对 目标集 （图 1 右边）作出的新预测。隐变量 $z$ 本质上是一个关于 $x$ 到 $y$ 之间映射（函数）的有限维嵌入。重要的是： $z$ 是一个随机变量，这使神经过程变成了一个概率模型，让我们能够捕获到函数的不确定性。一旦训练了此模型，就可以将 $z$ 的（近似）后验分布作为先验，并用于测试阶段的预测。

将数据集分为 背景集 和 目标集 ，乍一看很像标准的训练集和测试集，但事实并非如此，因为 目标集 也会被直接用于训练神经过程模型，并且我们的（概率）损失函数明确地定义在 目标集 上。这将使模型避免过拟合，实现更好的 样本外泛化。在实际工作中，我们会反复将训练数据分成随机选择的 背景集 和 目标集 ，以获得良好的泛化效果。

让我们考虑两种场景：

场景 1: 基于一个单一数据集，推断一个关于函数的分布。

场景 2: 当可以访问多个在某种程度上相关的数据集时，推断一个关于函数的分布。

Fig02

图 2: 两种场景设置

第一种场景对应于标准（概率）监督学习设置。给定某个由 $N$ 个样本构成的数据集，即给定 $(x_i, y_i)$ ，其中 $i = 1,\ldots ,N$ ，假设存在一个真实的函数 $f$ ，能够产生 $y_i = f(x_i)$ 的值，目标是学习函数 $f$ 的后验分布 $p(f | \cdot)$ ，并用它来获得测试点处的预测密度 $f(x^{\ast})$ 。

第二种情况可以从元学习的视角来看。给定 $D$ 个序号为 $d = 1,\ldots, D$ 的数据集，其中数据集由 $N_d$ 对样本点 $(x_i^{(d)}, y_i^{(d)})$ 组成，如果假设每个数据集都有自己底层产生 $y_i = f_d(x_i)$ 的函数 $f_d$ ，则我们可能想学习每一个函数 $f_d$ 的后验，并将其用于一个新的数据集 $d^*$ 。这种信息共享机制来源于一个基本假设，即：存在一个可以作为所有函数 $f_d$ 底层基础的共享过程。例如，在高斯过程的背景下，我们可以假设所有 $f_d \sim \mathcal{GP}$ 共享了相同的核超参数。如果学得了共享过程，则当给定一个新数据集 $d^{\ast}$ 时，我们就可以将关于函数的后验分布用作先验，进行基于函数样本的回归。

我想强调这两种场景的原因如下: 通常，在最标准的设置中，高斯过程回归是在第一种情况下进行的。即使在 $N$ 较小的情况下，这也倾向于工作得很好。但神经过程背后的思想似乎主要来自于元学习，在元学习设置中，隐变量 $z$ 可以被认为是一种在不同数据集之间共享信息的机制。无论如何，因为具有概率模型的元素，神经过程应该同时适用于这两种场景。

下面我们将分别研究神经过程在两种场景下的表现。

2 神经过程是如何实现的？

2.1 模型架构

这里有一个详细的神经过程生成模型的示意图。

Fig03

图 3: 神经过程的生成式模型架构图

让我们按照步骤来了解该生成式机制：

（1）每个背景点 $(x_c, y_c)$ 都通过神经网络 $h$ 被映射为隐表示向量 $r_c$ 。
（2）若干 $r_c$ 被聚合（实际工作中采用了逐元素求平均）在一起，以获得单一的隐表示向量 $r$ （一个与 $r_c$ 具有相同维度的向量）。
（3）向量 $r$ 被用于定义隐变量 $z$ 的概率分布，即 $p(z | x_{1:C}, y_{1:C}) = \mathcal{N}(\mu_z(r), \sigma_z^2(r))$ ， $z$ 被假设为高斯的。
（4）为了获得对目标点 $x_t^{\ast}$ 的预测，我们对 $z$ 进行采样，并将样本与 $x_t^{\ast}$ 拼接（concatenate），然后通过神经网络 $g$ ，将拼接向量 $(z, x_t^{\ast})$ 映射为 $y_t^{\ast}$ 的一个预测分布样本。

注：整个模型中，仅存在 $z$ 这一个随机向量。整个构造形式与变分自编码器极其相似。

2.2 参数估计和推断

（1）推断目标

神经过程需要实现如下推断目标：

对编码器 $h$ 和解码器 $g$ 的神经网络参数的点估计，如果聚合器 $a$ 也存在参数化形式，则同样需要对其参数做点估计；
隐变量 $z$ 的概率分布，鉴于高斯假设，实际上是实现对超参数 $\mu_z$ 和 $\sigma_z$ 的推断，具体实现方式可以考虑重参数化技巧。

（2）推断方法

与变分自编码器类似，神经过程的推断是在变分推断框架下进行的。具体来说，引入两个变分分布：

$q(z | \text{context})$ ：用于近似编码器网络输出的条件先验 $p(z | \text{context})$ （注：采用背景集单独训练得出）。
$q(z | \text{context, target})$ ：用于近似编码器网络输出的预测分布 $p(z | \text{context, target})$ （注：采用背景集和目标集一起训练得出）。

其中： $\text{contex} := (x_{1:C}, y_{1:C})$ ， $\text{target} := (x_{1:T}^{\ast}, y_{1:T}^{\ast})$ 。

近似分布 $q(z | \cdot)$ 被选择为具有如下推断模型图的具体形式。也就是说，我们使用相同的 $h$ 来映射 背景集 和 目标集 ，以获得聚合的 $r$ ，而这又被映射到 $\mu_z$ 和 $\sigma_z$ 。整个过程实现了 $q(z | \cdot) = \mathcal{N}(\mu_z, \sigma_z)$ 的参数化形式。

Fig04

图 4: 神经过程的推断模型图

（3）目标函数

变分推断的目标函数通常设定为变分下界，神经过程的变分下界为：

\text{ELBO} = \mathbb{E}_{q(z | \text {context}, \text {target})} \left[ \sum_{t=1}^T \log p(y_t^{\ast} | z, x_t^{\ast}) + \log \frac{q(z | \text {context})}{q(z | \text {context}, \text {target})} \right]

其中包含两个项:

第一项是 目标集的预期对数似然。需要对 $z\sim q(z | \text {context, target})$ 采样（ 图 4 左侧），然后使用获得的 $z$ 样本，结合 目标集 进行预测（ 图 4 右侧）。
第二项是 $q(z | \text{context, target})$ 和 $q(z | \text{context})$ 之间的负 $\mathbb{KL}$ 散度（ 图 4 左侧），具有正则化的效果。需要注意的是：其形式与 VAE 等常用设置略有不同，KL 散度形式为 $\mathbb{KL}(q || p)$ ，其中常用的目标分布 $p$ 在此时理应是先验 $p(z)$ 。本模型之所以不同于常规设置，是因为在模型中并没有直接定义 $p(z)$ ，而是定义了一个条件先验 $p(z |\text{context})$ 。由于该条件先验取决于编码器 $h$ ，我们无法获得其精确值，因此用近似分布 $q(z | \text {context})$ 来代替。

3 实验

第 3.1 节讨论将神经过程作为函数上的先验时，其行为表现，以便对 “函数上的分布” 这个概念有一个直观理解；第 3.2 节先不考虑待学习的函数类，仅简单地从几个样本点出发了解神经过程的训练方法；第 3.3 节将样本点范围扩展到某个简单函数类（如 $y = f_a(x) = a \sin(x)$ ，其中 $a$ 为该函数类的超参数 )的可能样本；第 3.4 节将样本点范围扩展到根据观测数据拟合的某个高斯过程类，模拟了一个完整的神经过程流程。

3.1 神经过程作为先验时的表现

首先，让我们探索一下 “将神经过程用作函数的先验分布时的” 行为表现，也就是说，在没有观测到任何数据、没有做任何训练之前的表现。

在初始化神经网络权重之后（从标准正态分布中随机采样），编码器 $h$ 和解码器 $g$ 均已确定，则我们可以对 $z \sim \mathcal{N}(0, I)$ 进行采样，并在覆盖输入域 $x^{\ast}$ 的网格上，通过对先验预测分布的采样，生成函数的随机样本（注：这是典型的贝叶斯先验预测分布的生成过程，常用于对先验分布进行评估，参见贝叶斯方法相关文献）。

相对于具有可解释核参数的高斯过程先验而言，神经过程先验的明确性要差得多，各种架构的选择结果（如使用多少个隐藏层，使用什么激活函数等）都会隐式地影响着其在函数空间上的先验分布。

例如，当使用 Sigmoid 激活函数并在 $\{1, 2, 4, 8\}$ 中选择 $z$ 的维度时，对神经过程先验做随机初始化得出的典型结果可能如 图 5 所示。

Fig05

图 5: 采用 Sigmoid 激活函数时，从神经过程先验中抽取的函数

而当使用 ReLU 激活函数来代替 Sigmoid 时，则将神经过程先验，则可能倾向于将概率质量放在不同的函数集上。

Fig06

图 6: 采用 ReLu 激活函数时，从神经过程先验中抽取的函数

3.2 在一个小数据集上训练神经过程

假设我们所拥有的是以下五个数据点。

Fig07

图 7: 5 个数据点构成的小数据集

神经过程的训练程序涉及到将数据集分割为 背景集 和 目标集 。可以选择固定大小的 背景集，也可以考虑选择同大小的 背景集 以覆盖更广泛的场景（例如，在每次迭代时，我们可以从 $\{1，2，3，4\}$ 的集合中随机抽取背景点数量）。

我们首先在这些随机子集上完成训练，然后就将训练好的模型（图 4 左侧）用作先验，并以所有数据为条件（即把所有五个点作为背景点）抽取并绘制后验函数样本。

下面的动画说明了在训练过程中，神经过程的预测分布变化情况。

Fig08

图 8: 在训练过程中，神经过程的预测分布的变化情况

因此，神经过程似乎已经成功地学习了一个穿过所有五个点的映射（函数）的概率分布。

现在让我们来探讨一下该神经过程的泛化情况，也就是说，我们在训练阶段和预测阶段采用不同的 背景集，看看会发生什么情况。下图是以红点为条件时得出的后验。显然，神经过程模型只在五个蓝点的子集上进行了训练，并不能泛化到不同的背景点集。因此，为了得到一个泛化性能更好的模型，可以考虑在一个更大的函数集（类）上对神经过程进行（预）训练。

Fig09

图 9: 不同背景集下做预测的泛化效果

3.3 在一个简单函数类上训练神经过程

到目前为止，我们已经探索了使用固定数据集进行训练的情况。但要想拥有一个泛化性能好的神经过程，我们似乎应该在一个大得多的 函数类 上训练它。，为此，本节先探讨一下在一个简单函数类中，神经过程应该如何表现。

让我们考虑一个玩具场景，我们观测的不是一个单一的函数，而是一小类函数。具体来说，让我们考虑所有形式为 $a\cdot\sin(x)$ 的函数，其中 $a \in [-1, 1]$ 。

我们感兴趣的是：

(1) 神经过程是否能够捕获到这个函数类？
(2) 神经过程是否能实现超越这个函数类的泛化？

让我们使用下面的训练过程：

(1) 定义一个函数样本：
- 从均匀分布 $a \sim U(-2, 2)$ 中抽取样本 $a$ 。
- 从均匀分布 $x \sim U(-3, 3)$ 中抽取样本 $x_i$ 。
(2) 定义该函数样本下的数据样本 $y_i := f(x_i)$ ，其中 $f(x)= a\sin(x)$
(3) 将 $(x_i, y_i)$ 随机分成 背景集 和 目标集 ，并执行一个优化步骤
(4) 重复

在这里， $z$ 使用二维向量形式，以便直观地查看模型学到的内容。在训练了神经过程之后，我们将对应于各种 $(z_1, z_2)$ 值的函数绘制在一个网格上可视化，如下图所示。

Fig10

从图中看来，从左到右的方向基本上是对我们的参数 $a$ 的编码。下面是同样效果的另一种可视化，我们改变其中一个潜在维度（ $z_1$ 或 $z_2$ ）。

Fig11

请注意，上面我们在预测时没有使用任何 背景集 ，而只是预先指定了各种 $(z_1, z_2)$ 值。现在让我们来看看使用这个已训练的模型进行预测的表现。

将 背景集 作为点 $(0, 0)$ ，如左图所示，将产生一个相当广泛的后验，看起来相当不错，涵盖了在 $a$ 的一定范围内类似于 $a\sin(x)$ 的函数（但注意，不是对所有 $a\in [-2, 2]$ 进行训练）。

Fig12

添加第二个背景点 $(1, \sin(1))$ 将产生中间所示的后验。与之前的图相比，后验发生了变化，例如， $a$ 为负值的函数不再被包括在内，但没有一个函数穿过给定的点。当增加遵循 $f(x)=1.0\sin(x)$ 的背景点的数量时，那么神经过程后验将变得合理地接近于真实的基础函数，如右图所示。

现在让我们来探讨一下，经过训练的神经过程在它所训练的那类函数之外的泛化程度如何。具体来说，我们来探讨一下它对以下函数的泛化情况： $2.5\sin(x)$ 和 $|\sin(x)|$ 。第一个需要从训练数据中进行一些推断。第二个函数的形状与训练集中的函数相似，但与其他函数不同，它的值是非负的。

Fig13

从图中可以看出，神经过程还不能超越它在训练期间所看到的东西。在这两种情况下，模型的行为在某种程度上是预期的（例如，在左边， $a=2$ 对应于模型所见函数类中的最佳拟合）。然而，请注意，在神经过程的预测分布中没有太多的不确定性。当然，过度自信的预测并不是神经过程所特有的，然而，与更多可解释的模型相比，它们的黑箱性质可能使它们更难诊断。

3.4 在来自高斯过程的函数上训练神经过程

根据到目前为止的实验，神经过程似乎不是高斯过程的开箱即用的替代物，因为后者在后验不确定性方面有更理想的特性。为了实现神经过程的类似行为，我们可以使用大量的高斯过程先验抽样来训练它。我们可以这样做。

(1) 从真实高斯过程中抽取函数样别 $f \sim \mathcal{GP}(0, k_{\theta}(\cdot))$ 。
(2) 生成函数样本对应的数据样本
- 抽取 $x_i \sim U(-3, 3)$
- 定义 $y_i := f(x_i)$
(3) 将 $(x_i, y_i)$ 随机分成 背景集 和 目标集 ，并执行优化步骤
(4) 重复进行

上述过程可以用固定的核超参数 $\theta$ 或不同的混合值来进行。例如，RBF（或平方指数）核有两个参数：一个控制方差（基本上是函数值的范围）；另一个控制 “扰动性”，被称为长度尺度参数。在这里，我们通过从高斯过程先验中抽取长度尺度值为 $\{1, 2, 3\}$ 的函数来说明其效果。

Fig14

为了涵盖神经过程训练中的各种函数，我们可以为 $\theta$ 指定一个先验 $p(\theta)$ 并从中抽取样本。在玩具实验中，我们在 $\{1, 2, 3\}$ 中均匀地改变长度尺度，并将 $z$ 的隐空间定义为二维，可以直观地看到神经过程学到了什么。

Fig15

这看起来非常酷！这个神经过程似乎已经学会了 $z$ 的二维隐空间，使我们可以在不同的函数之间平滑插值。

现在让我们探索一下使用这个神经过程得到的预测，比较其表现与高斯过程后验相比如何。我们采用逐渐增加数量的背景点集合（分别为 $\{3, 5, 11\}$ 个背景点），并考虑如下两个函数：

首先，使用一个相对平滑的函数 $f(x) = \sin(0.5x)$ ，预测结果如下。

Fig16

然后，让我们考虑 $f(x)==sin(1.5x)$ :

Fig17

在第一种情况下，神经过程的预测与观测结果相当接近，而在第二种情况下，有三个观测结果，它看起来很好，但当给它更多的点时，它还不能捕获到模式。这种影响很可能源于结构选择，即采用了过小的神经网络和低维的 $z$ 。为了改善模型的行为，使其更接近于高斯过程，我们可以考虑作出以下调整：

(1) 二维 $z$ 空间能够学习的东西有限，可以考虑使用更高维的 $z$ ，对于 $r$ 也同样如此。
(2) 在神经网络 $h$ 和 $g$ 中使用更多的隐藏单元，并考虑使其更深。
(3) 在训练阶段，观测更多函数样本以及更多的函数类（即高斯过程核超参数的更多变化）。

这种变化确实可以导致更理想的结果。例如，将 $\text{dim}(r)$ 增加到 $32$ 维， $\text {dim}(z)$ 增加到 $4$ 维，将 $h$ 和 $g$ 中增加更多的隐藏单元，我们得到了更好的表现。

Fig18

4 结论

尽管神经过程结合了神经网络和类似高斯过程的贝叶斯模型的元素来捕获函数的分布，但在从高斯过程到神经网络的整个谱段上，神经过程更接近于神经模型。通过仔细选择神经结构以及神经过程的训练程序，有可能实现理想的模型行为，例如类似高斯过程的预测不确定性。然而，这些影响大多是隐性的，这使得神经过程作为一种先验的解释更具挑战性。