Understanding the Distillation Process from Deep Generative Models to Tractable Probabilistic Circuits
理解从深度生成模型到易处理概率电路的提炼过程
https://proceedings.mlr.press/v202/liu23x/liu23x.pdf
摘要
概率电路(PCs)是一个通用且统一的计算框架,用于处理可解析的概率模型,这些模型能够高效地执行各种推理任务(例如,计算边缘概率)。为了使复杂现实世界任务具备此类推理能力,Liu等人(2022)提出从表达力更强但解析性较差的深度生成模型中提炼知识(通过潜在变量赋值)。然而,目前仍不清楚是哪些因素使得这种提炼能够取得良好效果。在本文中,我们从理论和实证两方面发现,概率电路的性能可以超过其教师模型。因此,我们不再仅仅从最具表达力的深度生成模型中进行知识提炼,而是研究教师模型和概率电路应具备哪些属性才能实现良好的提炼性能。这为我们带来了通用的算法改进以及针对特定数据类型的改进,相较于现有的潜在变量提炼流程。在实证方面,我们在具有挑战性的图像建模基准测试中大幅超越了当前最优技术水平的概率模型。特别是在ImageNet32数据集上,概率电路达到了4.06比特/维度,仅比变分扩散模型(Kingma等人,2021)落后0.34比特/维度。
引言
开发能够执行各种推理任务(例如计算边缘概率)的可解析概率模型(TPMs)至关重要,因为它们能够支持广泛的下游应用,如受限生成(Peharz等人,2020a;Correia等人,2020)、因果推理(Wang & Kwiatkowska,2022)和数据压缩(Liu等人,2021)。概率电路(PCs)(Choi等人,2020)是指一类具有相似表示形式的TPM,包括和积网络(Poon & Domingos,2011)、与或搜索空间(Marinescu & Dechter,2005)和算术电路(Darwiche,2002)。为了充分利用概率电路吸引人的推理特性,一个关键挑战是提高它们在复杂现实世界数据集上的建模性能。
近年来,从算法(Correia等人,2022;Shih等人,2021;Dang等人,2022;Peharz等人,2020b)和架构(Peharz等人,2020a;Dang等人,2021)两个角度对概率电路进行了大量的扩展和改进。特别是,Liu等人(2022)提出了潜在变量提炼(LVD)流程,该流程使用解析性较低但表达力更强的深度生成模型(DGMs)提供额外的监督,以克服基于期望最大化(EM)的概率电路参数学习器的次优性。通过潜在变量提炼,概率电路能够在ImageNet32(Deng等人,2009)等具有挑战性的数据集上实现与一些广泛使用的深度生成模型相当的性能。
然而,尽管潜在变量提炼具有巨大的潜力,但我们对其在何时以及如何导致更好的建模性能方面的理解仍然有限。因此,现有潜在变量提炼流程的成功在很大程度上依赖于反复试验。例如,如图1所示,当我们改变深度生成模型或从它们那里获得监督的策略时,即使概率电路的大小相似,建模性能也会发生显著变化。本工作旨在揭示潜在变量提炼(LVD)流程的奥秘,并为图像数据提供实用的设计指南。通过从变分推理的角度解释LVD,我们证明了LVD学习的概率电路(PCs)的性能并不必然受其教师深度生成模型(DGMs)的上限限制。这与从大神经网络向小神经网络提炼知识时,训练性能通常会下降的情况(Gou等人,2021)截然不同。因此,我们不应试图使用当前最优的深度生成模型来执行LVD,而应关注一个更基础的问题:教师深度生成模型的哪些属性会导致LVD学习的概率电路性能更好?尽管目前尚无确切答案,但我们确定了导致概率电路既具有表达力又紧凑的实用设计指南。
遵循这些指南,我们观察到现有LVD流程中存在一个普遍缺陷。具体来说,由于概率电路要求的离散潜在变量赋值与神经网络的连续表示之间存在不匹配,因此通常使用一次性离散化方法。然而,这会导致大量信息丢失,从而降低建模性能。为了解决这个问题,我们提出了一种渐进增长算法,利用来自概率电路的反馈进行动态聚类,从而最小化离散化造成的性能损失。渐进增长算法还能够利用可重用的子结构,从而生成紧凑且富有表达力的概率电路。结合从指南中得出的几个针对图像的特定设计选择,我们在三个具有挑战性的图像建模数据集上(CIFAR(Krizhevsky等人,2009)和两个下采样的ImageNet数据集(Deng等人,2009))大幅超越了当前最优的概率模型。特别是,我们在ImageNet32上实现了4.06比特/维度的性能,优于一些难以处理的流模型(Flows)和变分自编码器(VAEs),如Glow(Kingma & Dhariwal,2018),并且仅比当前最优的基于扩散的深度生成模型(Kingma等人,2021)高出0.34比特/维度。实验重现代码可在https://github.com/liebenxj/LVD-PG.git找到。
背景
本节介绍概率电路(第2.1节)和潜在变量提炼(LVD)流程(第2.2节)。
2.1 概率电路
概率电路(PCs)是一类广泛的概率图模型(TPMs),它们将概率分布表征为深层计算图。概率电路的语法和语义如下。
定义1(概率电路)。概率电路p(X)被表示为一个参数化的有向无环计算图(DAG),它通过一个单一的根节点nr定义了一组随机变量X上的联合分布。DAG中的节点被分为三种类型的计算单元:输入单元、和单元与积单元。值得注意的是,DAG中的每个叶节点都作为输入单元,而内部节点则根据其组合子分布的机制可以进一步细分为和单元或积单元。在前向路径中,每个内部节点从其子节点(用in(n)表示)接收输入,并计算输出,从而以递归方式编码一个概率分布pn。
概率图(PCs)能够准确且高效地回答众多概率查询(例如边缘分布、熵)(Vergari等人,2021),这一能力使其区别于各种深度生成模型。这种能力通常被解释为可处理性。
为了保证概率图的可处理性,必须对其有向无环图(DAG)结构施加一定的结构约束。例如,平滑性和可分解性确保概率图能够在线性时间内(相对于其大小,即其DAG中的边数)计算出任意边缘概率。这些是概率图单元n的变量范围φ(n)的属性,即包含其所有后代节点的变量集。
2.2 潜在变量蒸馏
尽管最近在开发高效的概率图(PCs)计算框架方面取得了突破(Dang等人,2021;Molina等人,2019),但利用大规模概率图的额外表达能力仍然极具挑战性。大量实证证据表明,这一现象的原因是现有的基于期望最大化(EM)的优化器无法在概率图庞大且层次嵌套的潜在空间中找到良好的局部最优解(Peharz等人,2016),这个潜在空间由其有向无环图(DAGs)中层次分布的求和单元定义。
潜在变量蒸馏(LVD)通过为概率图优化器提供额外的监督,即来自不太易处理但更具表达力的深度生成模型(DGMs)的语义感知潜在变量(LV)分配,克服了上述瓶颈(Liu等人,2022)。具体而言,LVD的操作步骤如下:首先,在概率图中实现一些/全部潜在变量(LVs)的具体化。也就是说,将原始的概率图p(X)转换为p(X, Z),其中X的边缘分布保持不变,即p(X) = ∑_z p(X, Z=z)。
接下来,使用深度生成模型(DGMs)为训练集中的每个样本诱导出语义感知的潜在变量Z的分配,从而得到一个增强的数据集。这一潜在变量诱导步骤可以通过多种方式和使用不同的深度生成模型来完成。例如,在Liu等人(2022)的研究中,Z是通过聚类由掩码自编码器(Masked Autoencoder,He等人,2022)产生的潜在特征来获得的。
最后,使用增强的数据集D_aug来最大化对数似然的下界,如最右边的项所示。
在使用增强数据集进行训练后,我们可以通过对潜在变量Z进行边缘化来获得目标分布p(X),这一过程可以在与概率图大小成线性的时间内完成(Choi等人,2020)。然后,可以使用原始数据集对概率图进行微调,以进一步提高性能。
潜在变量蒸馏(LVD)的成功主要归功于其能够简化概率图中深层嵌套潜在变量空间的规模和深度(Peharz等人,2016)。具体来说,在潜在变量具体化之后,深度生成模型(DGMs)可以为这些潜在变量提供监督,而基于期望最大化(EM)的概率图参数学习器则仅负责推断其余隐式定义的潜在变量的值。由于深度生成模型通过提供的潜在变量分配来指导概率图的学习,我们将其称为教师模型,将概率图称为学生模型。
潜在变量蒸馏中的性能差异特征
尽管潜在变量蒸馏(LVD)已经展现出其在提升大型概率图(PCs)性能方面的潜力,但其有效性在很大程度上取决于具体化潜在变量(LVs)的设计选择以及它们如何从外部来源进行诱导。具体来说,如图1所示,设计选择不当将导致性能显著下降,而设计得当则可以进一步缩小与当前最优但不易处理的深度生成模型(DGMs)之间的性能差距。因此,关于潜在变量蒸馏的一个关键但尚未回答的问题是:为了实现良好的建模性能,潜在变量诱导过程应遵循哪些设计原则?
我们通过变分推断(VI)来表征教师深度生成模型(DGM)和学生概率图(PC)之间的性能差异,从而为这个问题提供了一个初步的答案。变分推断是各种深度生成模型(如变分自编码器VAEs(Kingma & Welling, 2013)和扩散模型(Ho et al., 2020))的数学基础。考虑一个潜在变量模型pθ(x) := ∑_z pθ(x|z)pθ(z)。由于直接最大化对数似然log pθ(x)可能不可行,变分推断提出同时学习一个变分后验qφ(z|x),并最大化以下对数似然的证据下界(ELBO):
尽管形式不同,但这个ELBO目标与方程(1)在除去一个与概率图无关的常数因子后是等价的(详见附录A的严格阐述)。直观上,qφ(z|x)被视为外部模型,用于为每个训练样本x诱导潜在变量z的分配。因此,我们称方程(3)为潜在变量蒸馏(LVD)目标。
潜在变量蒸馏目标为表征教师深度生成模型(方程(2))的性能与学生概率图的对数似然之间的差异提供了桥梁。具体来说,如图2所示,教师深度生成模型与潜在变量蒸馏目标之间的性能差距,我们称之为潜在变量蒸馏差距,它表征了潜在变量蒸馏的性能损失。然而,教师深度生成模型与学生概率图之间的最终性能差异可能会远小于潜在变量蒸馏差距。具体来说,由于概率图的可处理性,ppc(z|x)可以以封闭形式获得。因此,在通过潜在变量蒸馏目标训练了ppc(x|z)和ppc(z)之后,我们可以“免费”地消除潜在变量蒸馏目标与概率图对数似然之间的变分差距。也就是说,如图2所示,在进行了潜在变量蒸馏训练(给出了底部的ELBO)之后,我们可以直接获得一个对数似然(显示在左上角)已知的概率图。
也许令人惊讶的是,上述分析表明,学生概率图(PC)的对数似然并不一定受教师深度生成模型(DGM)的证据下界(ELBO)的上限约束。具体来说,如图2所示,当变分差距大于潜在变量蒸馏差距时,就会出现这种情况。在极端情况下,如果学生完美模拟了教师(即pθ(x|z) = ppc(x|z)且pθ(z) = ppc(z)),则概率图就变成了的一个可处理的实例化,能够计算各种概率查询。为了实现这一理想情况,我们需要最小化潜在变量蒸馏差距。
缩小潜在变量蒸馏差距的一个关键见解是确保pθ(x|z)(以及pθ(z))与学生概率图ppc(x|z)(以及ppc(z))具有相似的建模假设。这在两个方向上都有很好的效果:一方面,通过模仿深度生成模型的归纳偏置(即pθ(x|z)和pθ(z)),我们可以学习到性能更好且参数更少的概率图;另一方面,尽管这可能会导致深度生成模型的性能下降,但去除深度生成模型中概率图由于其结构限制而无法完全采用的建模假设通常是有益的。接下来,我们首先确定建模假设不匹配的一般来源,并提出一种算法来缓解这个问题(第4节)。然后,我们展示这些设计原则如何专门应用于图像数据并提高建模性能(第5节)。
从连续神经表示中进行潜在变量蒸馏
最小化潜在变量蒸馏(LVD)差距的一个主要挑战是,具有表现力的连续神经嵌入与从概率图中具体化的潜在变量(LVs)的离散性质之间存在不匹配(因为求和单元表示的是离散混合)。因此,为了获得离散的潜在变量分配,要么使用事后离散化步骤(例如K-means),要么深度生成模型(DGM)需要学习离散表示,但这通常会导致性能比学习连续特征时差。这种一次性离散化策略会导致相对较大的潜在变量蒸馏差距,从而显著降低概率图(PC)的建模性能。
形式上,深度生成模型的变分后验可以分解为:
在这里,gφ 是一个神经网络,h 是一个连续的神经表示。尽管我们仍然需要对 h 进行离散化以获得概率图(PC)的潜在变量(LV)分配,但这个离散化过程不必是一次性的:我们可以利用概率图的反馈来调整并重新分配聚类索引,以缩小潜在变量蒸馏(LVD)差距。具体来说,我们希望聚类索引(即 z)的分配方式既能尊重神经表示 h(即具有相似 h 的样本被分配到同一个聚类中),又能被以聚类为条件的概率图轻松学习。直观上,虽然后者确保了能够被正确学习,但前者保证了 z 保留了 h 中的信息,这在实际应用中会导致一个更好的。
接下来,我们描述所提出的渐进增长算法,该算法克服了上述一次性离散化方法的次优性。该算法以数据集及其对应的连续神经嵌入作为输入,定义为。在确定了与范围 X 对应的潜在变量 Z 后,算法还接受一个初始的以聚类为条件的概率图 p(x|z) 作为输入。我们假设 Z 最初取单一值(即 D_train 中的所有样本都属于同一个聚类),因此 p(x|z) 由一个单头概率图表示。
给定一个预定义的聚类数量 K,渐进增长旨在学习一个离散化函数,该函数将每个 h 映射到一个聚类索引 i ∈ [K],以及一个 K 头概率图,表示 p(x|Z=i)(∀i ∈ [K])。直观地,一个 K 头概率图编码了在同一组变量上定义的 K 个分布,每个分布都根据定义 1 表示为一个子概率图。然而,与直接用 K 个独立的概率图编码分布不同,K 头概率图中的子概率图共享某些结构和参数。这是通过迭代地将 D 分割成更多的聚类,并相应地学习以聚类为条件的概率图的结构和参数来实现的。具体来说,如图 5 所示,渐进增长通过重复以下四个主要步骤来操作,详细步骤如下所述。
其中,φ 是概率图(PC)的参数。我们使用标准的小批量期望最大化(EM)算法(Peharz 等人,2020a;Choi 等人,2021)来优化方程(4)。超参数的详细信息见附录B。为了学习一个既紧凑又富有表现力的概率图,我们在参数学习阶段之后应用了Dang等人(2022)提出的剪枝算法。这导致概率图的规模显著减小,同时性能损失可忽略不计。
步骤2:重新分配聚类索引。由于一步离散化方法的次优性,λk分配的聚类索引可能无法完全尊重概率图p(x|z)。也就是说,由于λk是通过聚类神经表示h获得的,因此一些被分配到聚类i的样本x可能由在上一步中训练的pφ(x|Z=j)(j≠i)更好地建模。为了缓解这个问题,我们利用概率图的反馈来重新分配聚类索引。具体来说,如图5所示,样本x的聚类索引被重新标记为z := argmax_{i∈[k]} p(x|Z=i)。相应地,函数λk也被修改以反映这一变化。
在接下来的步骤中我们将详细阐述,这个重新标记的过程使我们能够摆脱过去迭代中分配不当的聚类,并且对渐进增长的有效性至关重要。
步骤3:选择生长的聚类。顾名思义,渐进增长是通过迭代增加Z中聚类数量来操作的。为了提高以聚类为条件的概率图(即方程4)的整体性能,我们选择平均对数似然较低的聚类进行进一步划分。具体来说,如图5所示,我们首先计算每个聚类i∈[k]的平均对数似然:
步骤4:扩展概率图和重新聚类数据。假设前一步选择了一组用于扩展的聚类索引I。这一步的目标是将这些|I|个聚类扩展为M个新聚类(M > |I|)。在内部实现中,我们需要对相应的样本子集进行重新聚类,并对概率图进行结构修改以适应新的聚类。以下描述了这两个过程。
为了确保多头概率图的结构和参数在重新聚类步骤后仍然与聚类分配λk相关,一种自然的方法是分别对概率图和数据集进行层次扩展和聚类。具体来说,对于每个选定的聚类i ∈ I,我们使用K-means将属于第i个聚类的样本聚类为n个聚类,并基于pφ(x|Z=i)为新增的聚类创建n-1个新的概率图根单元。我们采用了一种稍作修改的方法来同时对所有|I|个聚类的训练样本进行重新聚类。具体来说,我们首先选择属于I中聚类的样本子集(x, h)。然后,我们运行K-means对神经表示h进行聚类,将其分为K个聚类,其中前|I|个聚类中心的初始值设置为I中聚类的中心。之后,更新λk以反映新的聚类分配。这样,前|I|个新聚类仍然与相应的概率图pφ(x|Z=i)(i ∈ I)相关。
为了表示新添加的聚类,概率图的结构需要被修改,以增加M - |I|个额外的根/头单元来表示p(x|Z=i)(i ∈ {k+1, ..., k+M-|I|})。一个简单的策略是直接为新聚类复制M - |I|个现有根单元的所有后代单元。然而,这将显著增加以聚类为条件的概率图的规模,使得渐进增长算法的效率非常低。此外,这种方法排除了重用对不同z条件下建模x有用的子电路的可能性,严重限制了概率图在任何特定规模下的表达能力。
为了缓解这个问题,我们提出了一个结构增长算子,该算子仅复制描述分布时最重要的子结构。通过在原始和复制的子电路之间引入额外的边,概率图可以学习共享可用于描述各种z条件下的p(x|z)的结构。增长算法的核心是一个称为“流”的统计量,它衡量一个节点/边相对于样本x的生成重要性(Dang等人,2022;Liu & Van den Broeck,2021),定义如下。
直观上,函数 Fn(x) 量化了单元 n 对 x 的对数似然的“贡献”。图 4 展示了一个示例 PC-样本对,每个节点上方都标有似然值。相对流量高的节点和边被标记为红色。需要注意的是,高节点似然并不保证高流量,这一点通过 n1 和 n2 来说明:它们都有高似然值,但只有 n1 有高流量。对于数据集 D,Fn(D) := ∑x∈D Fn(x) 测量了 n 对 D 中样本的总贡献。
回想一下,我们的目标是将当前的 k 头 PC 扩展,使其具有 M − |I| 个额外的根单元,以分别编码 p(x|Z = i)(i ∈ {k+1, ... , k+M −|I|})。为了实现这一点,我们首先扩展了多头 PC 的定义 4。具体来说,内部节点的递归定义保持不变,对于第 i 个根节点,如果 x 被 λm 分配到聚类 i,则设置流量为 1,否则为 0。
在算法 1 中展示了提出的增长操作符。它由两个主要部分组成:在 3-4 行中,使用电路流量选择一组“重要”的(即流量高于预定义阈值的)节点进行生长;在 5-15 行中,以一种只复制选定节点的方式修改 PC,而其他部分保持不变。
回顾一下,我们的目标是将当前的k头概率图扩展为具有M - |I|个额外的根单元,以分别编码p(x|Z=i)(i ∈ {k+1, ..., k+M-|I|})。为了实现这一目标,我们首先扩展多头概率图的定义4。具体来说,虽然内部节点的递归定义保持不变,但对于第i个根节点,如果λm将x分配给聚类i,则我们将流设置为1,否则设置为0。
所提出的增长算子如算法1所示。它主要由两部分组成:在第3-4行中,利用电路流来选择一组“重要”(即流高于预定阈值的节点)节点进行增长;在第5-15行中,以仅复制选定节点而保持其他部分不变的方式修改概率图。在我们的用例中,由于我们想要修改与选定的|I|个聚类相对应的子电路,因此我们使用其聚类索引在I中的样本子集来调用算法1。根据流的定义,返回的概率图将有k + |I|个头,因为选定的|I|个根节点将被算法复制,而其他节点则不会。渐进增长在上述四个步骤之间交替进行,直到我们将聚类数量扩展到预定值K。因此,通过第4步增长的多头概率图的参数将在算法下一次迭代的第1步中更新。
综上所述,第4步中的数据重新聚类过程确保了聚类分配符合神经表示,而第2步中的聚类分配过程则导致了良好拟合的以聚类为条件的概率图。
缩小图像数据的LVD差距
以图像数据为例,本节展示了如何将第3节中介绍的缩小LVD差距的一般指导原则转化为实际的设计选择。在本文中,我们采用矢量量化变分自编码器(Vector Quantized Variational Autoencoders,简称VQ-VAEs)(van den Oord等人,2017;Razavi等人,2019)作为教师模型。接下来,我们首先简要介绍VQ-VAE,然后描述我们为更好地使pθ(x|z)(及pθ(z))的建模假设与ppc(x|z)(及ppc(z))的建模假设保持一致所做的设计选择。
如图6所示,VQ-VAE由编码器(生成特征图)和解码器(使用特征图重建输入图像)组成。与许多其他深度生成模型(Deep Generative Models,简称DGMs)不同,VQ-VAE的潜在特征图是由包含M个向量(代表M个码字)的码本构建的。具体来说,潜在特征图中每个位置的嵌入都必须是码本中的一个向量。由于特征图中的每个潜在码字都对应于输入图像的一个图像块,我们在潜在特征图的每个位置实例化一个LV Zi,并将相应的图像块定义为Xi。设Z := {Zi}i,图像x的似然度可以计算为p(x) = ∑z p(z) ∏i p(xi|zi)。对于每个Zi,我们可以使用上一节详细描述的渐进增长算法,从图像块xi和VQ-VAE生成的相应连续特征向量hi中学习p(xi|zi)。然后,使用生成的离散化函数可以为每个训练样本生成z = {zi}i,并用于训练p(z)。
然而,在上述处理中,VQ-VAE和概率图(Probabilistic Circuit,简称PC)所做的建模假设之间存在不匹配。首先,由于VQ-VAE生成的潜在特征图在所有位置都使用相同的码本,因此不同图像块的聚类条件分布应该是同质的。也就是说,对于每个离散化函数λ和样本对(x, h),我们有∀i, j ∈ [|Z|],(此处|Z|表示Z的集合大小,即潜在特征图中码字的数量,但原文此处的表述在此上下文中略显不完整,因为后文未直接给出关于i, j的具体性质或关系)。
为了反映这种归纳偏置(inductive bias),我们并不是为每个i ∈ [|Z|]独立地学习p(x|zi),而是将它们的训练样本进行聚合,并学习一个单一的聚类条件分布,然后将该分布应用于每个图像块。也就是说,我们在不同的聚类条件分布之间进行参数绑定(parameter tying)。这不仅可以减少概率图(Probabilistic Circuit,简称PC)的参数数量,还使我们能够使用更多的数据来训练更好的聚类条件分布。
另一个建模假设的不匹配来自于概率图(Probabilistic Circuit,简称PC)所假设的在给定Zi的条件下,Xi与X\Xi(即除了Xi之外的所有图像块)之间的条件独立性。然而,VQ-VAE的卷积解码器打破了这一假设,因为在给定zi的情况下,xi可能与其他图像块相关。为了缓解这种不匹配,我们使用了一个独立的解码器,其中zi是生成xi的唯一信息来源。虽然这会降低VQ-VAE的性能,但如第3节所示,通过局部变分蒸馏(Local Variational Distillation,简称LVD)学习到的PC的性能可以超越教师模型。LVD的主要目标是找到一组初始参数,这些参数可以通过期望最大化(Expectation-Maximization,简称EM)算法优化到良好的局部最优解。我们将在第6.1节中展示这一现象。
实验
本节首先实证验证第3节中的理论发现(第6.1节)。然后,我们转而评估我们的方法在图像建模基准上的表现(第6.2节)。
6.1 分析LVD中的性能差距
我们实证研究了第3节中的发现,即学生概率图(Probabilistic Circuit,简称PC)的对数似然可以超越教师深度生成模型(Deep Generative Model,简称DGM)的证据下界(Evidence Lower Bound,简称ELBO)。具体来说,我们考虑了局部变分蒸馏(Local Variational Distillation,简称LVD)流程的一个实例化,并实证计算了图2中所示的三个ELBO。为了简化计算,我们使用VQ-VAE作为教师模型,并采用一次性K均值离散化策略在Imagenet32数据集上训练一个PC。教师DGM的ELBO结果为-2493(图2右上角),而LVD目标的ELBO为-2499(图2底部)。因此,LVD差距为5,这与第4节中提到的极端情况相符(即学生几乎完美地模拟了教师)。因此,PC成为了教师DGM的一个可处理的实例化。得益于PC的可追溯性,我们能够免费地关闭变分差距,并获得一个对数似然为-2317的PC,这使得学生PC的表现优于教师DGM。
6.2 图像建模基准测试
我们在三个自然图像基准测试上评估了提出的局部变分蒸馏(Local Variational Distillation,简称LVD)流程算法改进:CIFAR(Krizhevsky等人,2009)和两个下采样的ImageNet(ImageNet32和ImageNet64)(Deng等人,2009)。
基线方法。我们将提出的方法(称为带渐进增长的LVD,即LVD-PG)与四个可追踪概率模型(Tractable Probabilistic Models,简称TPMs)基线进行比较:LVD(Liu & Van den Broeck,2021)、隐藏Chow-Liu树(Hidden Chow-Liu Tree,简称HCLT)(Liu & Van den Broeck,2021)、随机和积网络(Random Sum-Product Network,简称RAT-SPN)(Peharz等人,2020b)和爱因斯坦求和约定网络(Einsum Network,简称EiNet)。这些基线涵盖了最近大多数关于扩展和提高TPMs表达能力的努力。此外,为了评估与不太易处理的深度生成模型(Deep Generative Models,简称DGMs)之间的性能差距,我们还进一步将LVD-PG与以下基于流的模型和变分自编码器进行了比较:Glow(Kingma & Dhariwal,2018)、RealNVP(Dinh等人,2016)、BIVA(Maaløe等人,2019)和变分扩散模型(Variational Diffusion Models,简称VDM)(Kingma等人,2021)。基线方法的实现细节见附录C。
实证见解。我们首先从与其他TPM模型的性能比较开始。如图7所示,在所有基准测试中,LVD-PG始终大幅优于以前的方法。特别是,在CIFAR上,一个约1200万参数的LVD-PG模型明显优于一个通过LVD训练的约8亿参数的PC;在ImageNet32上,通过LVD-PG训练的约2000万参数的PC也比通过原始LVD训练的约8亿参数的PC获得了显著的性能提升。这表明,适当的设计选择可以进一步挖掘LVD训练表达力强且紧凑的PC的潜力,从而显著提升大型PC的性能。
接下来,我们将LVD-PG实现的性能与三个采用的DGM基线进行比较。值得注意的是,如表1所示,我们的方法使PC能够在ImageNet64上超越除当前最优VDM之外的所有DGM,而在ImageNet32上,LVD-PG仅比BIVA差0.1位每像素差分(bit per dimension,简称bpd),比VDM差0.34 bpd。
消融研究。为了评估第4节中提出的渐进增长算法和第5节中详细阐述的针对图像数据的特定修改(如在VQ-VAE中使用“独立解码器”)的效果,我们分别训练了两个没有这些组件的PC,并进行了消融分析。这两个PC的模型大小与在ImageNet32上用LVD-PG训练的当前最优PC(约2.6亿参数)相似。具体来说,与具有4.06位每像素差分(bit per dimension,简称bpd)的当前最优LVD学习PC相比,没有渐进增长的LVD学习PC仅达到4.12 bpd,而使用卷积解码器的LVD学习PC的性能降至4.18 bpd。
相关工作
近年来,为了扩展和提高概率电路(Probabilistic Circuits,简称PCs)的表达能力,人们做出了巨大的努力。许多工作侧重于构建既具有表达力又紧凑的初始PC结构(Rahman等人,2014;Adel等人,2015;Rooshenas & Lowd,2014),而其他工作则旨在实现一个逐步增加模型容量的迭代结构学习过程(Di Mauro等人,2021;Dang等人,2020;Liang等人,2017)。这些方法在各种密度估计数据集(如MNIST系列数据集)上取得了显著的性能提升。
然而,仅改进PC结构似乎并不能在自然图像和文本等现实世界的高维数据集上带来太多的性能提升。为了解决这个问题,最近有许多研究致力于探索将PC与神经网络(Neural Networks,简称NNs)相结合的不同方式,以获得既易处理又具有表达力的混合模型。例如,条件SPN(Shao等人,2022)利用NNs的表达力来学习具有表达力的条件密度估计器;HyperSPN(Shih等人,2021)使用NNs来正则化PC的参数;Correia等人(2022)在NNs表示的连续潜在空间模型的帮助下学习PC的连续混合。
上述在扩展PC方面取得成功的关键是开发了计算框架和易于使用的库,这些框架和库使大规模PC的训练变得高度高效。具体来说,EiNet(Peharz等人,2020a)和SPFlow(Molina等人,2019)利用成熟的深度学习包(如PyTorch,Paszke等人,2019)来实现各种推理和参数学习过程,而Juice.jl(Dang等人,2021)则实现了自定义内核以更好地处理稀疏PC。
结论
本文旨在揭示从难以处理的深度生成模型(Deep Generative Models,简称DGMs)到易于处理的概率电路(Probabilistic Circuits,简称PCs)的潜在变量蒸馏过程。我们发现了理论和实证证据,表明学生PC的性能可以超过教师DGM的性能,这种性能提升源于PC的易处理性,它能够“免费”地缩小变分差距。基于这种对蒸馏技术的变分解释,我们进一步提出了算法改进,从而在性能上显著超越了当前最优的树形概率模型(Tree-structured Probabilistic Models,尽管原文为SoTA TPMs,但根据上下文我推测是指某种最优的树形或层次结构的概率模型,因为TPM通常不直接指代最优模型,这里为了表述清晰我进行了适当的调整)。此外,该方法还优于多个难以处理的DGM基线模型。
https://proceedings.mlr.press/v202/liu23x/liu23x.pdf