Scaling Continuous Latent Variable Models asProbabilistic Integral Circuits

连续隐变量模型的概率积分电路扩展

https://arxiv.org/abs/2406.06494v2

摘要

概率积分电路（Probabilistic Integral Circuits, PICs）最近被提出作为一种概率模型，它具备表达力强的生成模型背后的关键要素：连续潜在变量（LVs）。PIC 是符号化的计算图，将连续潜在变量模型定义为由函数组成的层级结构，这些函数通过相加、相乘，或对某些潜在变量进行积分来组合。如果潜在变量可以解析地积分掉，那么 PIC 是可追踪的；否则，可以用一种称为 QPC（Quadrature PC）的概率电路来近似，该电路编码了一种分层数值积分过程。到目前为止，只有树状结构的 PIC 被研究过，而通过数值积分训练它们在大规模数据上需要大量内存。本文中我们解决了这些问题，并提出了：(i) 一个基于任意变量分解构建有向无环图（DAG）结构 PIC 的流程；(ii) 一个使用张量化电路架构训练 PIC 的方法；以及 (iii) 神经函数共享技术以实现可扩展训练。在广泛的实验中，我们展示了函数共享的有效性以及 QPC 相比传统 PC 的优越性。

1 引言

连续潜在变量（LVs）可以说是许多成功生成模型的关键要素，从变分自编码器 [24] 到生成对抗网络 [20]，再到最近的扩散模型 [56]。虽然这些模型能够从数据中学习复杂的分布，但它们通常仅限于采样，在执行诸如计算边缘分布或条件分布等基本概率推理任务时并不高效，甚至不可行。另一方面，对于（分层）离散潜在变量模型来说，这些任务可能是可追踪的 [4,3]，但在大规模数据上学习这些模型更具挑战性 [9,10,32,33]。

这种在可追踪性、学习难度和表达能力之间的权衡，可以通过一种新提出的深度生成模型——概率积分电路（PICs）[18] 来分析和探索。PIC 使用符号函数电路定义连续潜在变量的层次结构。当其连续潜在变量可以解析地积分掉时，PIC 是可追踪的。而对于不可追踪的 PIC，也可以系统性地用（张量化）概率电路（PCs）[53,4] 进行近似，后者是离散潜在变量模型的表示语言。这类 PC 实例编码了用于近似 PIC 的分层数值积分过程，因此被称为“积分 PC”（QPC）。

从 PIC 中提取 QPC 已被证明是一种有效的训练 PC 的替代方式，但目前仅针对树状结构的 PIC 进行了探索。如何构建并扩展到更丰富的潜在变量结构仍然是一个开放的研究问题，需要新的工具 [18]。在本文中，我们填补了这一空白，通过重新定义 PIC 的语义，并将其扩展到 DAG 形式的连续潜在变量层次结构。具体来说，我们将 PIC 设计为一种表示分层拟张量分解的语言 [49]，并通过轻量级多层感知机进行参数化。

贡献如下：

  1. 我们提出了一个系统的流程，可以从任意变量分解出发构建 DAG 结构的 PIC（第 3.1 节）。

  2. 我们展示了如何通过一种分层积分过程来学习和近似 PIC，并将其编码为与已有电路架构 [43,42,31,36] 匹配的张量化 QPC（第 3.2 节）。

  3. 我们提出了函数共享技术以实现 PIC 训练的可扩展性，使我们可以用多头多层感知机（MLP）来参数化 PIC，其资源需求与 PC 相当（第 3.3 节）。

  4. 在广泛的实验中（第 4 节），我们展示了：(i) 函数共享在扩展性方面表现出色；(ii) 与通过 EM 或 SGD 常规训练的 PC 相比，QPC 表现更优，而且它们是从参数数量少达 99% 的 PIC 中提炼出来的。

2 概率积分电路

符号说明。我们将输入变量表示为X，将潜在变量（LVs）表示为Y和Z，它们的取值分别记作x、y和z。我们用小写字母表示标量（例如，w ∈ ℝ），用粗体小写字母表示向量（例如，），用粗体大写字母表示矩阵（不包括 X、Y、Z，例如 ），用粗体花体字母表示张量（例如，。

一个概率积分电路（Probabilistic Integral Circuit, PIC）c是一个符号化的计算图，它表示关于观测变量X和连续潜在变量Z的非负函数c(X) = ∫ c(X, z) dz，即对于所有x，有c(x) ≥ 0。与概率电路（Probabilistic Circuits, PCs）[53, 4] 类似，PIC 也包含输入单元、求和单元和乘积单元。然而，与 PC 不同的是，PIC 操作的是函数而非标量，并且使用了一种新型的单元：积分单元 ∫，它可以用来表示连续潜在变量。为了满足非负性，我们使用非负函数对 PIC 进行参数化，并采用正的求和权重，具体如下：

图 1(b) 和图 2(b) 展示了概率积分电路（PIC, Probabilistic Integral Circuits）的示例。请注意，我们使用f来表示附加在输入单元和积分单元上的函数，这些函数本质上是模型的参数；而我们使用g来表示所有类型单元所输出的函数。一个 PIC c c的输出是其根单元 u u返回的输出函数，该函数仅在输入空间 X X上定义，即 c ( X ) = g u ( X ) c(X)=gu(X)，并且 Z u = ∅ Zu=∅。

与概率电路（PCs）类似，在 PIC 上施加结构性约束可以实现可处理的推理 [4]。因此，我们假设如下结构性条件：
(i) 所有加法单元（+ 单元）接收的函数都定义在相同的输入变量上（即“平滑性”）；
(ii) 所有乘法单元（× 单元）接收的函数都定义在不相交的输入变量集合上（即“可分解性”）。

当一个 PIC 的局部变量（LVs）可以被解析地积分出去时，该 PIC 被认为是可处理的。这意味着我们可以穿过积分单元，计算它们所定义的积分问题，最终输出一个函数。
值得注意的是，当局部变量之间具有线性-高斯关系[27] 或当这些函数是多项式时，这种积分是可行的。

对于不可处理的 PIC，我们可以通过一种分层数值求积过程来进行近似，该过程可以被编码成一种称为求积概率电路（Quadrature PC，QPC）的 PC。直观上，PIC 中的每一个积分单元都可以用 QPC 中的一组加法单元来近似，这些加法单元以一些预先计算的求积值为条件，并包含大量但有限个输入单元 [18]。将一个 QPC 具体实现（materialize）出来，可以用于通过近似最大似然方法来训练 PIC：给定一个 PIC，其附加在输入、加法和积分单元上的参数的梯度，可以通过对应的 QPC 进行反向传播 [18]。这种方法也为训练 PC 提供了一种替代方案，其效果可以媲美传统的学习方法。

到目前为止，PIC 的构建仅限于从具有连续潜变量（LVs）的概率图模型（PGMs）[27] 进行简单的编译过程 [18]。简而言之，PGM 中的潜变量节点成为 PIC 模型中的积分单元，而 PGM 中的（条件）分布则成为与 PIC 的输入和积分单元相连的函数 fu，如我们在图 1 中所示。然而，为了避免分层求积过程导致 QPC 呈指数级增长从而阻碍学习过程，PGM 的结构必须被限制为一棵树。这使得当前的 PIC 具有了作为简单潜变量树模型 [3] 的语义，也明显限制了其表达能力，因为更复杂的潜变量之间的交互无法实现。要构建更具表达能力的 PIC，就需要重新诠释这一语义并引入新的工具，我们将在下文中完成这一点。

3 构建、学习与扩展 PICs

在第 3.1 节中，我们系统化了 DAG 形状 PIC 的构建过程，展示了如何从任意变量分解出发来构建它们，从而超越当前最先进的方法 [18]。随后，在第 3.2 节中，我们展示了如何使用编码分层求积过程的 QPC 来学习和近似此类 PIC，并复现了先前工作中提出的 PC 架构 [36]。最后，在第 3.3 节中，我们提出了（神经）函数共享技术，用以参数化 PIC，使得其 QPC 实例化过程快速且高效，从而实现更大模型和更大数据集上的扩展。

3.1 从任意变量分解构建 PICs

标准的概率电路（PCs）可以根据已建立的流程进行构建，这些流程能够灵活地表示任意变量分解以及丰富的离散潜变量（LV）交互 [36, 41]。在下文中，我们推导出一个类似的用于构建 PIC 的流程，使其能够处理连续 LV，从而超越当前仅限于树状语义的 PIC（见第 2 节），同时仍能在不使实例化的 QPC 规模爆炸的前提下执行分层求积。为此，我们首先形式化“层次变量分解”或“区域图”的概念，基于此我们将构建我们的 PIC 结构。

定义 1（区域图（RG）[15]）：
一个关于输入变量 X 的区域图 R 是一个二分的、有根的有向无环图（DAG），其节点要么是表示 X 子集的区域（regions），要么是指定某个区域如何被划分为其他区域的划分（partitions）（见图 2(a)）。

区域图可以（i）从概率图模型（PGMs）中编译而来 [5, 27, 3, 31]，（ii）随机初始化 [43, 16]，（iii）从数据中学习得到 [15, 19, 40, 57]，或者（iv）根据数据模态（如图像）构建 [45, 42, 36]。如果我们像图 1 那样从一棵 PGM 树中编译出来，那么生成的 RG 也将是一棵树，从而产生一个类似树结构的 PIC [18]。

我们在算法 1 中详细描述的流程接受任意 DAG 形状的 RG 作为输入，并能输出类似 DAG 的 PIC。不失一般性，我们假设有一个 RG R，它只允许（i）对区域进行二元划分（即所有乘积单元将有两个输入单元），以及（ii）单变量叶子节点，如图 2(a) 所示。

3.2 通过张量化 QPC 学习 PIC

给定一个 DAG 形状（难以处理的）PIC，我们现在展示如何使用编码分层求积过程的张量化 PC 来近似它，即 QPC。直观上，我们将 PIC 解释为编码一组准张量 [49]，这是一种张量的推广形式，其每个维度可能对应于连续潜变量（LV），从而在每个维度中具有潜在无限的条目。我们通过数值求积将其具体化为经典的张量。

我们从张量化电路的定义开始，并简要回顾数值求积的相关知识

从 PIC 到 QPC 给定一个候选的 PIC，我们将以后序遍历的方式对其进行遍历，并迭代地将每个 PIC 单元与一个电路层（定义 2）相关联，这一过程我们称之为实例化（materialization）。
我们在算法 3 中详细描述了这一过程，其本质上是在 PIC 单元上层次化地应用公式 (1)

每个单元编码了一个可能作用于连续和离散变量上的函数，因此表示一个准张量（quasi-tensor），我们通过仅在求积点上对该函数进行评估来近似它，从而将其实例化为一个经典张量（如图 3(a,b) 所示）。
为了便于进行求积计算，我们假设所有 PIC 的潜变量（LV）都具有有界的定义域Δ:=[−1,1]。
这样一来，我们就可以对每一个所需的（多变量）近似使用相同的求积规则 (z~,w~)，从而简化处理和表述。

3.3 利用神经函数共享扩展 PIC

QPC 的实例化过程可能会非常耗费内存和时间，这取决于以下三个因素：
(i) 我们需要实例化的函数的计算代价，
(ii) 所需函数评估的并行化程度，
(iii) 积分点的数量 K。

为了解决这些问题，我们引入了神经函数共享（neural functional sharing）[47]，即通过共享多层感知机（MLP），一次性参数化多个 PIC 单元。
这种方法使我们能够扩展到更大的模型和数据集，相比之前的工作 [18]，我们的实例化过程更快、更节省内存。

PIC 函数共享

函数共享之于 PIC，就如同参数共享之于 PC。
这种类型的共享可以应用于一组输入/积分单元——这些单元根据某些标准被分组在一起，且它们的函数都具有相同数量的输入变量和输出变量。

最后，我们介绍并应用两种不同的单元分组方式：

  1. 第一种方式是将所有输入单元分组在一起，这种方法只有在所有输入变量具有相同定义域时才适用。
    在这种分组下，结合 F-sharing，我们只需要实例化 K×P 个参数，并用它们来参数化每一个 QPC 输入层。

  2. 第二种方式是将所有处于 PIC 结构中相同深度的积分单元分组在一起，我们将这种方式与 C-sharing 相结合，并将其实例化为一个折叠的求和-乘积层（folded sum-product layer）。
    尽管将那些会实例化为折叠层的单元进行分组是一个自然且方便的选择，但请注意，我们也可以对那些不会以这种方式实例化的单元进行分组。
    一旦一个 PIC 中的所有单元都被分组，就可以按组进行实例化操作。

使用（多头）MLP 的 PIC 函数共享

与 [18] 类似，我们使用轻量级的多层感知机（MLPs）来参数化 PIC 的输入单元和积分单元。
不过，我们不再为每个函数单独使用一个 MLP，而是应用函数共享，以期使 QPC 的实例化过程更加快速且节省内存。

快速且内存高效的 QPC 实例化 结合PIC 函数共享按组实例化，使得通过数值求积训练 PIC 成为可扩展的方法，因为我们大幅减少了 QPC 实例化所需函数评估的数量及其计算代价。我们现在可以实例化非常大的 QPC，其规模可以与近期过参数化的 PC 模型相当，但在使用较大 K 的情况下，所需的可训练参数最多可减少 99%。这在原始的 PIC 公式 [18] 中是无法实现的，原因有两点：(i) 整个 QPC 是一次性（one-shot）实例化的，而不是按组进行的；(ii) 没有实现函数共享，因为每个输入/积分函数都有自己的 MLP。

4 实验

在我们的实验中，首先我们评估了函数共享在通过数值积分扩展 PIC 训练方面的有效性，并将其与标准 PC 以及不使用函数共享的 PIC [18] 进行比较。接着，借鉴以往的研究 [10, 32, 33, 18]，我们将 QPC 和 PC 作为分布估计器，在多个图像数据集上进行对比。我们始终使用梯形积分规则进行训练。所有实验都在 NVIDIA A100 40GB 显卡上完成。我们的代码可在 github.com/gengala/tenpics 获取。

得益于我们的流程，我们现在可以使用两种最近为图像数据设计的 RG（区域图），它们相比传统 RG 构建的结构具有更好的扩展性 [45, 43, 31]：

  • 四叉树

    （QTs）是树状的 RG，

  • 四图

    （QGs）是 DAG 形状的 RG [36]。

这些 RG 是完全平衡的，因此使用算法 1 对其进行构建时，会得到平衡的 PIC 结构，适合对积分单元进行按深度方向的 C-sharing（组合共享）。我们在附录 B 中详细介绍了 QT 和 QG 的结构。

我们将一个张量化架构表示为[RG]-[求和-乘积层]-[K]，例如QT-CP-16，它可以像标准 PC 一样训练，也可以从 PIC 实例化为 QPC。我们将像素视为分类变量，因此我们的架构建模的是概率质量函数。

扩展 PIC 对于每种模型类型{PC, PIC}，我们指定一个二元组(·, ·)，其中第一个（第二个）参数表示输入（内部）层/组所使用的共享技术，包括{F, C, N}，其中 N 表示不共享。

在图 5 中，我们报告了使用 PC ( )、使用带 ( ) 和不带 ( ) 积分单元组函数 C-sharing 的 PIC 进行一次 Adam [23] 优化步骤所需的时间和 GPU 内存。

我们注意到，使用函数共享的 PIC ( ) 在扩展方面非常有效，所需的资源与标准 PC ( ) 相当；而未使用函数共享的 PIC ( ) ——如先前的工作 [18]——则慢了几个数量级，并且在 K>64 时很快出现内存不足（OOM）。

值得注意的是，某些 PIC 的QG-TK配置 ( )，见表 D.2，所需的 GPU 内存甚至比 PC 更少，这是因为可训练参数的数量存在显著差异，因为 Adam 在优化过程中需要存储这些参数的副本。

实际上，标准 PC 和无函数共享的 PIC ( , ) 的参数数量达到了数亿级别（最高达 20 亿+），而使用函数共享的 PIC ( ) 则扩展得更为平滑，最多仅达到 600 万+ 参数。

我们强调，PIC 的可训练参数数量与我们实例化时的 K 无关，只与用于参数化的 MLP 大小 M 有关，这也可以理解为评估 PIC 函数的成本。

更多（表格形式的）细节请参见附录 D.1。

分布估计 借鉴以往研究 [10, 32, 33, 18]，我们广泛测试了 QPC 和 PC 作为分布估计器在标准图像数据集上的表现。

我们的完整结果见附录 D.2，此处我们仅报告表现最佳模型的每维比特数（bits-per-dimension, bpd），这些模型均属于QG-CP架构，反映了 DAG 形状 RG 具有更强的表达能力。

我们的完整结果还表明，我们为 PIC 引入的更具表达力但计算代价更高的 Tucker 层在 K 较小时表现最好，但在扩展时较为困难。

所有 QPC 都是从应用了输入单元 F-sharing 和积分单元组 C-sharing 的 PIC 实例化而来，即 PIC(F, C)。

我们首先以 MNIST 家族数据集开始，它包含 6 个灰度 28x28 图像数据集：MNIST [29]、FASHIONMNIST [55]，以及带有 4 个划分的 EMNIST [7]。

图 6 显示，QPC 总体表现最佳，在 6 个任务中有 5 个优于标准 PC，在 6 个任务中有 5 个优于基于复杂启发式学习的 PC 方法（如剪枝与生长），并在全部 6 个任务中优于一些深度生成模型（DGMs）。

接下来，我们转向更大的 RGB 图像数据集：CIFAR [28]、ImageNet32、ImageNet64 [14] 和 CelebA [34]。

为了与之前的工作 [32, 33] 进行比较，我们需要使用YCoCg 变换对数据集进行预处理，这是一种有损的颜色编码方式，能持续提升 PC 在 RGB 图像上的性能。我们还报告了使用无损 YCoCg-R 变换 [39] 预处理的数据集结果，有效地将数据集数量翻倍。

关于这些变换的详细信息请参见附录 C.3。

从表 1 中可以看出，QPC 再次表现出很强的竞争力，一致优于使用 Adam 常规训练的标准 PC，也优于文献中表现最好的 PC 模型 HCLT（后者通过 EM 方案和分块方法训练 [32]）。

此外，QPC 的表现接近于通过潜变量蒸馏（LVD，表 1 中 LVD-PG）训练的 PC [32, 33]，这是一个需要额外监督其潜空间的框架，通过从现有深度生成模型（DGMs）中提取信息实现。

该技术需要预训练的 DGM、多种启发式方法，以及通过 EM 或 SGD 进行最终微调阶段，而我们的 PIC 训练方法则是端到端且自包含的。

5 讨论与结论

在本工作中，我们系统化了 PIC 的构建过程，将其扩展到了 DAG 形状结构（第 3.1 节），进行了张量化处理（第 3.2 节），并使用函数共享技术扩展了其训练能力（第 3.3 节）。

在我们的实验中（第 4 节），我们展示了当将 PIC 的可追踪近似模型 QPC 用作分布估计器时，这一流程表现得非常有效。这反过来也成为一种新的、有效的用于大规模学习 PC 的工具。

事实上，之前的研究已经表明，通过 EM 或梯度上升方法直接训练大型 PC 是具有挑战性的，而且随着模型规模的增加，PC 的性能会趋于饱和 [10, 8, 32, 33]。

我们的贡献超越了这些限制，同时提供了一种简单、原理清晰且完全可微的流程，其性能可以媲美更复杂的替代方案 [10, 32]（见第 4 节）。我们推测，这种效果的产生是因为通过 PIC 来训练 PC 极大地缩小了搜索空间，同时允许：
(i) 更平稳的训练动态，以及
(ii) 实例化任意大的可追踪模型。

可追踪模型的发展是机器学习中的一个重要任务，因为它们提供了多种推理机制，并可用于许多下游应用，例如表格数据建模 [8]、生成建模 [42]、无损压缩 [30]、遗传学 [11]、知识图谱 [37]、受限文本生成 [58] 等等。

我们的工作也与张量网络和**(准) 张量分解** [51, 26, 22, 6, 49] 有一定关联，正如 [36] 最近所展示的那样，使用张量化电路的语言可以表示分层张量分解。

此外，我们注意到最近提出的非单调 PC[35]（即具有负求和参数的 PC）也可以看作是从某些函数可能返回负值的 PIC 中进行求积过程的结果。

我们的工作并非没有局限性。尽管我们展示了使用函数共享训练 PIC 所需的资源与标准 PC 相当，但像 VAE、流模型（flows）和扩散模型（diffusion models）这类传统的连续潜变量模型仍然更具扩展性。

此外，目前还无法从 PIC 中进行采样，因为我们无法对我们的（多头）MLP 进行可微采样。

未来的工作可能包括探索更高效的 PIC 训练方式，例如使用 LVD 或变分推断 [24] 技术来直接最大化 PIC 的下界，在最后阶段仅使用数值求积作为微调步骤，以蒸馏出高性能的可追踪模型。

我们相信，我们的工作将在生成建模领域，特别是可追踪模型的研究方向上，激发新的研究进展。

B 区域图（Region Graphs）

在算法 B.1 中，我们详细描述了四叉树（Quad-Tree，QT）和四图（Quad-Graph，QG）区域图 [36] 的构建过程。
具体来说，将输入标志isTree设置为True（对应 QT）或False（对应 QG）即可构建相应的区域图。

直观上，这些区域图通过递归地将图像分割成小块（patches），直到每个区域恰好对应一个像素为止。
分割操作同时在水平方向和垂直方向进行，后续的小块可以共享（从而生成非树状的区域图，即 QG），也可以不共享（生成树状区域图，即 QT）。

C.3 YCoCg 颜色编码变换

在图 C.1 和图 C.2 中，我们提供了实验中所使用的无损和有损版本 YCoCg 变换的 PyTorch 代码（见第 4 节）。在图 C.3 中，我们展示了如何应用这些变换，并表明有损版本的平均误差小于一个比特。最后，在图 C.4 中，我们展示了将这两种变换应用于 RGB 图像时，在视觉效果上的显著差异。

D.2 附加的分布估计结果

在本节中，我们报告了所有实验的表格结果。

请注意，标准 PC 的每个输入层都由一个 K×P 的矩阵进行参数化，其中 P 是类别数量：对于灰度图像数据集为 256，对于 RGB 图像则为 768（即 256×3）。

我们发现，在所有像素之间共享一个单一的输入层，会使灰度图像的性能（略有）下降（详见表 D.3）。
相比之下，我们发现这种共享方式对 RGB 图像数据集的性能有明显提升。

除了提升 RGB 图像数据集的性能之外，这种共享方式还将可训练参数的数量从 D×K×P 显著减少到仅 K×P，其中 D 是像素数量。

例如，如果我们为一个用于 64x64 图像、K=256 的张量化架构的所有输入层分别参数化，则需要

256×64×64×768=805,306,368

个参数；而如果应用共享，则只需 3,145,728 个参数。

因此，如果不采用这种共享方式，我们在 GPU 上甚至无法扩展到较大的张量化架构（如 QG-CP-512）。

所有 QPC 都是从 PIC 实例化而来，其中对输入单元组使用了F-sharing（全共享），对积分单元组使用了C-sharing（组合共享）。

我们广泛地将 QPC 和 PC 作为密度估计器在多个图像数据集上进行了比较，并在表 D.3、表 D.4 和表 D.5 中报告了测试集上的每维比特数（bits-per-dimension, bpd）。

原文链接： https://arxiv.org/abs/2406.06494v2