分离容量、目标函数、批量大小、估计器和步长对流变分推断（flow VI）的影响

Disentangling impact of capacity, objective, batchsize, estimators, andstep-size on flow VI

分离容量、目标函数、批量大小、估计器和步长对流变分推断（flow VI）的影响

https://arxiv.org/pdf/2412.08824?

摘要

基于归一化流的变分推断（flow VI）是一种有前景的近似推断方法，但其在不同研究中的表现仍不一致。许多算法选择会影响 flow VI 的性能。我们进行了逐步分析，以厘清一些关键因素的影响：容量（capacity）、目标函数（objectives）、梯度估计器（gradient estimators）、梯度估计的数量（batchsize）以及步长（step-sizes）。每一步都单独考察一个因素，同时利用前一步骤的见解或通过大量并行计算来中和其他因素的影响。为了实现高保真评估，我们整理了一个合成目标分布的基准测试集，这些目标分布代表了常见的后验病理情况，并允许进行精确采样。我们针对各个因素提供了具体的建议，并提出了一种 flow VI 的实施方案（recipe），其性能达到或超过了当前主流的即插即用哈密顿蒙特卡洛（HMC）方法的表现。

1 引言

基于归一化流的变分推断（flow VI）[80] 通过一系列可逆变换构建灵活的变分族，以近似后验分布 [96, 3, 5, 38, 31] [73, 55]。尽管这一方法具有前景，但其表现仍不一致——一些研究报告了成功案例 [3, 5, 38, 91, 31]，而另一些则强调优化困难和较差的结果 [43, 10, 8, 24, 47, 60, 6]。

多个因素影响着 flow VI 的性能，使得难以确定不一致性的原因：是容量不足、目标函数不合适、梯度方差过高，还是步长设置错误？这些因素之间的“纠缠”限制了我们对 flow VI 能力的理解和最佳实践的形成，阻碍了其更广泛的应用。

本文采用逐步分析的方法，分别探讨几个关键因素的影响——容量（第3节）、目标函数（第4节）、梯度估计器与批量大小（第5节），以及步长（第6节）。每一步都回答一个关于某个因素的基本问题，同时通过应用前一步骤的见解或利用现代 GPU 集群的计算能力来中和其他因素的影响。例如，在第4节中我们提出问题：是否需要复杂的覆盖多模态的目标函数？为了测试这一点，我们通过借用第3节中验证过的流模型结构来中和容量的影响，并通过使用极大的批量大小和广泛的超参数扫描来中和优化选择的影响。

这种方法要求有精确的准确性衡量标准。然而，由于真实后验通常是难以处理的，这使得高保真评估变得复杂：在没有真实样本的情况下，适用于变分推断（VI）和哈密顿蒙特卡洛（HMC）的通用评估指标非常有限——证据下界及其变体 [16, 25, 29, 30] 仅适用于 VI，而像有效样本数量这样的收敛性度量则主要用于 HMC。虽然可以使用已有方法生成的样本作为真实样本的代理 [61]，但通用且可靠的评估指标仍然有限——Wasserstein 距离 [93] 虽常被使用，但其计算成本随样本数量增长较快 [21]，不适合用于高保真评估（见第2.2节）。

为了解决评估难题，我们整理了一组合成目标分布基准（第2.1节），这些目标反映了常见的后验病理特征：条件数差、非线性曲率、重尾分布以及参数间的相互依赖关系。这些目标提供了真实的样本，从而确保了高保真评估（评估策略详见第2.3节）。此外，我们使用了一个可扩展的 Wasserstein 距离替代指标——边缘-Wasserstein 距离（公式2），它通过对各个一维边缘分布的 Wasserstein 距离进行平均得到。该距离度量在样本数量上呈对数线性增长，能够高效而准确地比较 VI 和 HMC 方法（第2.2节）。

我们在下面展示了主要的研究发现。

容量（Capacity） 。Real-NVP 流 [27] 可以准确地逼近具有挑战性的目标分布（图3）。尽管 Real-NVP 被广泛使用 [96, 3, 38, 99]，但一些研究也报告了其表现不佳的情况 [10, 24, 47, 6]。我们通过中和其他因素的影响后表明，只要网络结构设计得当，Real-NVP 具备足够的容量来准确表示复杂的后验分布（第3节）。

目标函数（Objectives） 。在具备高容量流模型的前提下，复杂的目标函数是不必要的（图4）。一些工作建议使用覆盖多模态的目标函数，因为它们能更好地覆盖整个分布 [59, 95, 68]，但也有一些研究表明这些目标难以优化 [35, 37]。我们直接使用真实样本优化覆盖多模态的 KL 散度（KL(p ∥ q)），并验证这种方法确实可以取得很好的效果（图4）。然而，当模型容量足够时，更易于优化的寻找模态的 KL 散度（KL(q ∥ p)）已经足够有效（图4）。

梯度批量大小与估计器（Gradient batchsize and estimators） 。大批量大小显著提升了高容量流模型的表现（图5）。已有研究探索了用于流模型的低方差梯度估计器 [3, 91, 92]。虽然这些方法有所帮助，但我们发现仅靠它们还不够，更大的批量大小能够显著提升高容量流模型的性能，因此建议在可能的情况下将两者结合使用（第5节）。

步长与优化（Step-size and optimization） 。保持步长在一个狭窄范围内对收敛至关重要（图6）。自动选择合适步长仍然是一个开放性问题 [56, 97]。我们展示出即使在使用高容量流、低梯度方差和自适应优化器 [51] 的情况下，优化过程也可能在看似稳定运行数千次迭代之后突然发散（第6节）。实验表明，将步长控制在 之间，可以在不同目标上实现稳定的长期训练结果（图6）。

总体来看，有两个关键因素至关重要：高容量的流模型 （解决表达能力限制）和大的梯度批量大小 （简化优化过程）。有了这两个要素，一个简单的实施方案就已足够：使用标准的变分推断目标函数、尽可能使用降低方差的梯度估计器，并采用固定在小范围内的步长；合理初始化并参数化流模型架构，并使用像 Adam 这样的自适应优化器进行长时间优化 [51]（详见附录B）。

基于这一方案，我们展示了 flow VI 在具有挑战性的目标分布上，仅需较少的模型评估次数即可达到或超越主流 HMC 方法的表现（见第7节以及图1和图7）。

2 设定（Setup）

给定一个模型 p(z,y)，其中 z 是潜在变量（latent variables），y 是观测变量（observed variables），推断的目标是近似后验分布 p(z∣y)。变分推断（Variational Inference）通过最大化证据下界（ELBO）来学习一个近似分布 q [86, 49, 94, 12]，其中：

2.1 目标分布（Targets）

我们设计了一组合成目标分布的基准测试，用以捕捉现实世界后验分布中常见的病理特征（参见附录 L 获取目标细节）。所有这些目标分布都允许精确采样 ，使得原本无法进行的分析成为可能（见第3节和第4节）。精确样本还使得对 VI 和 HMC 方法的公平比较 成为可能（评估策略详见第2.3节）。图2展示了这些三维目标分布的样本。

2.2 评估指标（Metrics）

有一些性能评估指标仅适用于 VI 或 HMC 中的一种，而非两者通用：例如 ELBO、证据上界 [48] 和重要性加权 ELBO [16] 等标准指标并不适用于 HMC；而像有效样本数量这样的收敛诊断指标则不适用于 VI。

还有一些指标使得精确比较变得困难：预测测试似然度由于测试数据集的不确定性 [2] 而不可靠；而对于像自由度为 1.5 的学生 t 分布（Student-t）这类重尾目标分布来说，矩甚至可能不存在。

Wasserstein 距离 [93] 的优势在于它可用于对 VI 和 HMC 进行公平比较，但它需要参考样本，并且在样本数量增加时计算效率较低 [21]。在初步实验中，我们发现即使使用了数千个样本，Wasserstein 距离估计中的噪声仍然使得精细比较难以实现（参见附录 C 中图8）。

幸运的是，在单变量（一维）情况下，计算 Wasserstein 距离可以简化为排序操作。基于这一点，我们采用了边缘-Wasserstein （marginal-Wasserstein）度量——即各对应的一维边缘分布之间的 Wasserstein 距离的平均值。

设 A 和 B 为两个 S×d 的样本矩阵，其中 S 是样本数量，d 是维度数量。令 表示对 A 按列排序后的版本，使得 是第 j 列中的最小值。

与 Wasserstein 距离一样，当样本量相对较小时，这种度量也会受到蒙特卡洛噪声的影响，但由于它具有更好的可扩展性（scalability），我们可以使用更大的样本量来降低噪声（见附录 C）。

当然，边缘-Wasserstein 度量仅关注边缘分布，因此它无法捕捉维度之间的相关性。然而，由于其良好的可扩展性，它成为了高保真评估中唯一可靠的选择。

2.3 评估策略

为了进行可靠的性能评估，我们使用边缘-Wasserstein 度量（公式2）。即使在精确推断的情况下，该度量的最小可实现值也依赖于样本数量和目标分布的几何特性 [21]。为了体现理想性能的参考标准，我们绘制了从目标分布中独立抽取的两组样本之间的边缘-Wasserstein 度量，并用黑色虚线表示（如图1所示）。当一种推断方法的结果接近这条线时，说明其生成的样本质量与目标分布的样本相当。

对于所有合成目标的评估，我们都使用同一组一万个参考样本（见图1和图7）。在第7节中，我们还在一些真实目标上进行了评估，并通过长时间运行 HMC 方法生成参考样本（详见附录K）。

2.4 实验细节

为了模拟不同的规模，我们使用了三个维度设置：2维、10维和100维（除了 Funana 目标，它至少需要3个维度）。我们将 flow-VI 在 JAX [14] 中实现，并使用 TensorFlow Probability [58] 和 NumPyro [76] 中的实现来进行 HMC 方法的实验。所有方法均在 Nvidia A100 GPU 上运行。如需完整细节，请参见附录 G 到 K。

3 Real-NVP 流是否具备足够的容量来表示具有挑战性的目标分布？

Real-NVP 流是推断研究者中一种广受欢迎的选择 [96, 3, 89, 24, 85, 38, 87, 99, 6]。然而，也有一些研究报告了其表现不稳定或效果不佳的情况 [24, 47, 6]。目前尚不清楚这些结果是由于 Real-NVP 流在表达能力上的不足，还是优化过程中的失败所导致。

为了解决这一问题，我们在中和其他因素的前提下 ，研究容量（capacity）的影响，并提出如下问题：Real-NVP 流是否具备足够的容量来准确表示具有挑战性的目标分布？

为了回答上述问题，我们直接优化 KL(p ∥ q)。具体来说，我们使用来自目标分布的大批量样本进行优化，从而大大简化了训练过程。（实际中用于优化 KL(p ∥ q) 的算法通常会受到基于自归一化重要性采样的高方差梯度的影响 [35, 37]，但在使用精确样本的情况下，这些问题可以被规避。）

此外，优化 KL(p ∥ q) 通常是更优的选择，因为由此得到的近似分布相较于通过优化标准变分推断目标函数 KL(q ∥ p) 得到的近似结果，预期能够更好地覆盖整个目标分布 [67, 68, 71, 50]。

我们通过扫描不同的步长和学习率调度策略、并进行大量迭代优化，来中和其他优化因素的影响（详见附录 G）。

图3 展示了当我们通过增加耦合层的数量或神经网络中的隐藏单元数量来提升模型容量时，边缘-Wasserstein 度量的变化情况（详见附录 E）。关键发现是：只要合理选择这些网络结构参数，Real-NVP 流具备足够的容量来准确表示复杂的后验分布 。

唯一的例外是自由度为 1.5 的学生 t 分布（Student-t），在这种情况下，KL(p ∥ q) 的优化因重尾特性而遇到困难 [47]（相比之下，KL(q ∥ p) 表现良好，见图4）。

主要发现包括以下几点：

1. 层数 ：增加层数非常有效。当层数达到10层或更多时，性能往往能够接近精确推断的结果（随着层数增加，曲线逐渐靠近黑色虚线）。

2. 隐藏单元数 ：增加隐藏单元对高维问题更有帮助（红色曲线与蓝色曲线之间的差异随维度增加而增大）。

尽管人们通常认为增加容量会提升表达能力，但这些结果表明，现有的流模型（如 Real-NVP）已经具备表示复杂后验几何结构的能力。因此，只要做出适当的优化选择，即使是“相对简单”的流模型，也有望取得令人印象深刻的结果。

4 Flow-VI 是否需要覆盖多模态的目标函数？

前一节表明，Real-NVP 流具备足够的容量来表示具有挑战性的目标分布。然而，第3节依赖于使用精确样本 进行优化，这使得我们有必要重新考虑所使用的优化目标。

选择合适的优化目标是一个活跃的研究领域。一些研究者提倡使用“覆盖多模态”的目标函数（mode-spanning objectives），例如第3节中使用的 KL(p ∥ q)，因为它们能够更好地覆盖整个目标分布 [59, 95, 25, 42, 64, 19, 101]。尽管这些方法前景良好，但其实用算法可能面临梯度方差大的问题 [78, 32, 37, 35]。

一个主要的替代方案是使用“更简单”的标准变分推断目标函数。然而，也有研究表明，在与流模型结合使用时，这种方法也可能遇到困难 [10, 24, 6]。

总体而言，目标函数的选择仍不明确，引发了一些关键问题：
我们是否真的需要复杂的覆盖多模态的目标函数？
当模型容量足够高时，标准的 VI 目标是否已经足够有效？

为了回答上述问题，我们在中和其他因素影响的前提下 ，专注于优化“更容易”的 KL(q ∥ p) 目标，并检验它是否足以实现我们在第3节中通过优化 KL(p ∥ q) 所观察到的良好性能。

为此，我们采用了大批量大小 （以减少方差），并对梯度估计器、步长和学习率调度策略 进行了详尽搜索，并进行了大量迭代（详见附录 H）。

图4 展示了在优化 KL(q ∥ p) 和 KL(p ∥ q) 后，边缘-Wasserstein 度量随耦合层数量变化的趋势。关键发现如下：

1. 目标函数 ：优化复杂但“覆盖多模态”的 KL(p ∥ q)（蓝色曲线）确实比优化“寻找模态”的 KL(q ∥ p)（红色曲线）取得了更好的结果，除了 Student-t 分布（自由度 ν = 1.5）的情况 ，这可能是由于其重尾特性导致的 [47]。

2. 容量 ：随着流模型容量的增加，两种目标之间的性能差距逐渐缩小，两者都能达到接近精确推断的结果。（仅在 Funana 目标的一百维情况下仍有明显差距，表明该问题具有非常复杂的几何结构。）

这些结果揭示了 flow VI 与使用简单变分族（如高斯分布）的传统 VI 之间的一个重要区别：
与其使用复杂的散度（divergence）来弥补目标分布与变分族之间的不匹配，不如直接增加流模型的容量并优化标准目标函数 。只要容量足够，散度形式的选择就不再那么关键 。

5 减小方差的估计器是否有帮助？

前几节表明，如果流模型具备足够的容量，优化标准的 KL(q ∥ p) 目标函数就可以逼近具有挑战性的目标分布。然而，第4节所采用的策略依赖于不切实际的超参数穷举搜索 。为了实现有效的实际应用，我们需要一种更高效的流程。为此，我们首先提出一个问题：如何估计梯度以实现可靠的优化？

优化 ELBO（公式1）的一个主要挑战在于它依赖于随机梯度估计 。已有不少研究探索了构建低方差梯度估计器的方法 [79, 82, 36, 66, 83, 90, 9, 33, 100, 17]。一个流行的选择是“sticking-the-landing”（STL）[83, 90] 估计器，它通常表现良好 [3, 23, 91, 92, 6]。然而，STL 要求对流变换 T 进行求逆 [3, 91, 92]，这会带来额外的计算成本和潜在的数值问题 [10]。对于像自回归流这样求逆代价高昂的流模型来说，STL 也不太实用 [53, 72]。

从原则上讲，现代 GPU 可以为梯度方差问题提供一种蛮力解决方案——只需并行抽取大量样本进行估计并取平均值。这引发了一些自然的问题：在使用高容量流模型时，梯度估计器的选择与批量大小之间是如何相互作用的？

为了回答这个问题，我们使用不同批量大小 ，分别采用 STL 和标准的总梯度估计器来优化一个高容量流模型。我们对不同的迭代次数进行独立优化，并对步长和学习率调度策略进行详尽搜索，以减少优化因素的影响（详见附录 I）。

图5 绘制了在十维目标下，边缘-Wasserstein 度量随梯度批量大小变化的趋势。关键发现如下：

1. 估计器选择 ：在小批量情况下，STL（红色）始终优于标准的总梯度估计器（蓝色），但随着批量大小增加，两者之间的差异逐渐缩小，并最终都达到了精确样本的精度水平。

2. 批量大小 ：对于任意固定的迭代次数（同一行数据），性能随着批量大小的增大而显著提升，突出了减小梯度方差的重要性。

这些发现强调了即使对于那些能够表示复杂后验分布的高容量流模型而言，减小梯度方差对于实证性能仍然至关重要。大批量大小可以显著提高性能，在硬件条件允许的情况下几乎都应该被使用。此外，STL 也能有效减小方差、提升收敛性，在可行的情况下也应被优先考虑（参见附录 F 获取效率分析）。

6 优化过程能否可靠地实现自动化？

虽然前一节关注的是可靠地估计 ELBO 的梯度 ，但其余的超参数搜索仍然不切实际。这引出了下一个重要问题：我们如何可靠地进行优化？

自动化变分推断（VI）的优化仍然是一个开放性问题 [56, 4, 5, 84, 97]。诸如优化器的选择、步长（step-size）、学习率调度策略（step-schedule）以及迭代次数的设定等决策并不直观 [56, 3]，而现有文献中也缺乏针对 flow VI 的具体实用指南。

已有若干使用流模型的研究报告称，Adam 优化器 [51] 取得了成功。即便如此，在使用 Adam 的前提下，选择最优的步长调度策略仍然是一个挑战。

这引发了一系列问题：

• 是否存在一个 适用于不同目标分布 的高容量流模型的通用 步长范围 ？

• 是否应该根据 维度大小对步长进行缩放 ？

• 能否在 几百次迭代内的早期表现 基础上预测出最佳步长？

为了透明地回答这些问题 ，我们使用十个不同的步长值 ，对高容量流模型进行了独立优化 ，并尝试了不同数量的迭代次数。我们采用了大批量大小 和 STL 估计器 （详见附录 J）。

图6 展示了边缘-Wasserstein 度量随步长变化的情况。从中可以得出一些关键结论：

1. 导致发散的步长 ：某些步长在运行数千次迭代时表现良好，但在更长时间运行后会出现发散。例如，在一百维 Funnel 目标中，一个步长为的设置在 1K 次迭代时表现良好，但在 10K 次迭代时却出现了发散。

2. 稳定的步长范围 ：步长在 范围内时，在多个目标分布上均表现出稳定良好的性能（一个显著例外是一百维的 Funnel 分布，此时稍大的步长反而效果更好）。

这些结果表明，仅凭早期几百次迭代的表现来预测最佳步长是具有挑战性的 。当使用梯度方差较低的高容量流模型时，在至少优化 10K 次迭代的前提下，将步长控制在 这一狭窄范围内，可以在不同目标分布和维度下实现最优性能（参见图中的紫色、灰色和黄色曲线）。

7 Flow VI 与 HMC 方法的比较如何？

至此，前几节的内容为我们提供了一个简单的实施方案（recipe）：

• 使用高容量流模型，

• 采用传统的 VI 目标函数（以简化优化），

• 在可行的情况下使用 STL 估计器（以降低梯度方差），

• 使用大批量大小（进一步降低方差），

• 以及在小范围内使用固定步长（详见附录 B 中的详细建议）。

自然而然地，我们提出一个问题：这一方案的实际效果究竟如何？

在本节中，我们将所提出的方案与当前最先进的即插即用式（turnkey）HMC 方法进行比较。

在现代 GPU 上使用 flow VI 是相对直接的——只需在每次迭代中并行计算更多随机梯度估计值即可。然而，要将 GPU 有效用于 HMC 方法则更具挑战性 [63, 45, 46]。对于像 NUTS 这样领先的即插即用式 HMC 方法来说，其复杂之处在于链依赖（chain-dependent）的控制流程使得并行化困难 [57, 75, 77, 45]。而对于一些最新方法，如 CHEES [45] 或 MEADS [46]，高效的跨链通信需求也使实现变得复杂。目前，运行多个并行采样链 是将 GPU 与 HMC 方法结合使用的最简单且最有效的方式 [62, 45]。

在实际问题中，计算瓶颈通常是模型评估 。因此，我们从两个维度来统计模型评估次数：

• 并行评估次数 （代表加速器规模的代理指标），
• 顺序评估次数 （代表运行时间的代理指标）。

对于 HMC 方法来说，并行评估次数对应于链的数量，顺序评估次数对应于 leapfrog 步数。

对于 VI 方法来说，并行评估次数对应于梯度批量大小，顺序评估次数对应于优化迭代次数。

由于当前框架的限制，我们对预热阶段（warmup phase）中的 leapfrog 步数进行了外推处理，参考了预热后阶段的数据（见附录 K）。

图7 绘制了边缘-Wasserstein 度量随顺序评估次数变化的趋势，针对不同的目标密度。除了合成目标外，我们还在六个现实世界问题上进行了比较（模型细节见附录 L）。关键发现如下：

1. 并行评估次数 ：随着并行计算资源的增加（从上到下穿越各行），所有方法的性能都有显著提升。

2. 顺序评估次数 ：在较小的批量大小下，flow VI 的表现与 HMC 方法相当或更优；而在批量更大时，flow VI 所需的顺序评估次数要少得多。

这些结果表明，flow VI 能够更有效地利用大批量大小，避免了基于链式采样的 HMC 方法中固有的协调开销。当然，未来 HMC 在现代加速器上的进一步发展可能会削弱这一优势；但就目前而言，使用我们提出的方案的 flow VI 是一个强有力的替代选择。

8 结论

本文通过逐步分析的方法，揭示了影响 flow VI 性能的关键因素。我们的分析发现，高容量流模型 和大梯度批量大小对于实现优异性能至关重要。

我们为目标函数、梯度估计器以及优化策略的选择提供了实用建议。此外，我们还展示了 flow VI 在具有挑战性的目标分布上，能够达到或超越主流即插即用式 HMC 方法的表现 ，且所需的顺序计算步骤要少得多。

A 相关工作（Related Works）

已有若干研究探索将归一化流应用于黑盒变分推断 （BBVI）[96, 8, 3, 24, 5, 6]。这些研究通常对流模型的能力做出一些先验假设，并更侧重于具体应用。

我们则专注于通过逐步分析的方法 ，厘清影响 flow VI 的关键因素。我们的研究结果应有助于未来此类方法的应用。

许多研究致力于自动化 BBVI 的某些方面 [56, 3, 23, 4, 5, 97]，旨在为概率模型提供无需人工干预的即插即用式解决方案。与这些方法不同，我们的工作重点在于理解不同因素的影响，提供详细的指导建议，而非追求完全自动化的解决方案。我们认为我们的分析为未来基于流模型的自动推断工具奠定了必要的基础。

Agrawal 等人 [3] 提出了一种结合归一化流、STL [83]、步长搜索方案以及后处理重要性采样步骤的方法，以提升现成的 BBVI 性能。但他们并未剖析影响 flow VI 的各个因素，使用了相对较小的批量大小，运行在 CPU 上，并且未对其失败案例进行深入分析。相比之下，我们的工作专门探讨了不同因素的影响，充分利用现代 GPU 进行大批量训练，并展示了适当优化后的 flow VI 可以达到或超越 HMC 方法的表现。

Andrade [6] 研究了如何在高维问题上稳定地优化 Real-NVP 流，并考虑了网络结构选择对优化稳定性的影响。我们的研究独立采用了与其最优选择一致的结构，从而避免了不稳定性。虽然他们提出的一些技巧可能进一步提升我们的性能，但我们的目标是全面理解各因素的影响，而不仅仅局限于结构设计。

Dhaka 等人 [24] 使用基于重要性采样的诊断方法评估性能，推荐使用归一化流，但由于优化困难而报告了较差的表现（见其附录中的图 C.2 和 C.3）。根据我们的研究发现，我们推测他们所遇到的问题可能部分源于相对较低容量的流模型（导致表达能力不足）以及较小的批量大小（导致梯度方差较高）。

Jaini 等人 [47] 揭示了具有重尾的目标分布需要近似分布具备相同尾部特性的基础分布。我们在图3中的发现支持了这一观点：优化 KL(p ∥ q) 在重尾目标上的表现不佳。更重要的是，在第4节中，使用标准 VI 目标函数和高容量流显著提升了性能。

Blessing 等人 [13] 强调了在推断研究中标准化评估的必要性，并提出了一个包含合成与现实世界问题的基准测试集，特别关注多模态目标。我们也使用了合成密度来构建受控环境，并采用如 Wasserstein 距离等积分指标进行评估。此外，我们还提供了一个理想性能的参考指标（图中黑色虚线，见第2.3节），以便更好地解读数值结果。我们还引入了一组六个现实世界问题，以展示 flow VI 的实际性能。

最近的文献也探索了将流模型作为 MCMC 方法提议分布的用途 [74, 44, 98, 68, 7, 65, 34, 40, 15, 50, 85, 18]，其中一些方法使用 KL(q ∥ p) 优化来初始化提议分布。这些方法与直接使用 flow VI 是正交的，并可能从本研究所揭示的见解中受益（例如学习更好的初始提议分布）。

B 推荐方案（Suggested Recipe）

要实现优异的 Flow VI 性能，需要进行几项关键的选择。然而，其中一些选择取决于可用的计算资源。为了便于更广泛地采用，我们提出一个推荐方案，并鼓励实践者根据自身约束条件进行适当调整。

绝对必需项

通常，在使用性能更强的 GPU 时更容易满足以下要求：

• 容量（Capacity） ：使用高容量的流模型以减少表示能力上的限制。
  我们默认使用至少包含十个耦合层、每层有32个隐藏单元的 Real-NVP 流模型（除非明确研究容量的影响），详见图3和图4，以及附录 E 中的结构细节。

• 批量大小（Batchsize） ：使用大量的梯度估计值（即大 batchsize）以降低梯度方差并简化优化过程。
  在非显式变化批量大小的实验中，我们使用了 个样本。不过，即使使用较小的批量大小也能取得不错的效果（见图5）。

在具备高容量流模型和大批量大小的前提下，以下建议表现非常出色：

• 目标函数（Objective） ：优化传统的 VI 目标函数。尽管覆盖多模态的目标函数有一定帮助，但标准目标更容易优化，并能达到相当的性能（见图4）。

• 梯度估计器（Estimator） ：使用“sticking-the-landing”（STL）梯度估计器以降低梯度方差。
  对于某些流模型来说，使用 STL 可能具有挑战性，请参见附录 F 的相关讨论。

• 步长（Step-size） ：选择一个固定在 范围内的步长（见图6）。如果资源允许，可在该范围内进行搜索。

• 优化策略（Optimization） ：使用 Adam 等自适应优化器进行大量迭代优化。
  实验表明，至少进行一万次更新在不同目标和维度下都表现出色（见图6）。

基于初步实验，我们还提出以下附加建议：

• 基础分布（Base distribution）

  ：在优化前评估不同基础分布在 ELBO 上的表现，并选择初始 ELBO 最高的那个。
  为了实验透明，我们默认使用标准正态分布作为基础分布。但在多个目标分布上，将拉普拉斯近似设为基础分布有助于优化的热启动 [29, 30, 3]。将学生 t 分布设为基础分布也被证明在重尾目标上具有优异的表现 [47, 60, 6]。我们建议在优化前进行小规模尝试，选择 ELBO 更高的那个作为基础分布。

• 参数初始化（Parameter initialization） ：将神经网络参数初始化为较小的值，等价于将流变换初始化为恒等变换 [3]。我们建议采用这种方式，以便更好地控制初始分布。

• 尺度函数的非线性选择（Non-linearity for scale function） ：在仿射变换中的尺度函数部分，我们建议使用双曲正切函数（hyperbolic tangent）（参见公式3及附录 E 中的相关讨论）。尺度函数所使用的非线性形式会影响仿射耦合流的稳定性 [10, 6]。我们尝试了多种非线性函数后发现，将双曲正切函数嵌套在指数函数中（如公式3所示）能够在稳定性和表达能力之间取得良好的平衡，并在多个目标分布上表现出一致的良好性能。

C n-Wasserstein 评估细节

两个分布之间的 Wasserstein 距离通常无法以闭合形式表示 [93]。用于计算它的数值方法在样本数量增加时扩展性较差 [21]。当使用有限数量的样本时，这些方法会产生噪声较大的估计值，使得 n-Wasserstein 距离在精细比较中不可靠。为了说明这一点，我们提供了一个简单的代表性示例。

C.1 Wasserstein 计算示例

考虑一个十维的标准正态目标分布。我们将计算这个目标分布与某些近似分布之间的 2-Wasserstein 距离和边缘-Wasserstein 度量。对于高斯分布来说，这两种 Wasserstein 度量都可以用闭合形式表示 [28, 20]。我们将把这些精确值 与通过数值方法计算出的值进行对比，以理解这些度量在精细比较中的有效性。

由于边缘-Wasserstein 度量具有更好的可扩展性，我们可以使用更大的样本量，并在相对较短的时间内获得极其精确的测量结果（见图8b）。

我们使用 ott-jax [22] 来计算 2-Wasserstein 距离，并在 Nvidia A100 GPU 上运行实验。

总体来看，图8 表明：除非使用大量的样本，否则具有不同程度准确性的近似分布可能会被评估为相同的错误 Wasserstein 值。2-Wasserstein 度量不适用于精细比较 ，因为它无法扩展到此类精确评估所需的大量样本。相比之下，边缘-Wasserstein 度量具有更强的扩展性，在合理时间内即可提供准确的测量结果 。

D 局限性（Limitations）

我们在实验中使用了合成目标分布 。这一选择是有意为之，目的是减少可能的变量干扰，并便于与 HMC 方法进行高保真度的性能比较。未来的研究可以深入探讨这些发现对现实世界问题 的影响。然而，我们预计我们的结论仍然适用。事实上，我们也展示了 flow VI 在多个现实问题上能够达到或超越 HMC 方法的表现（见图7和图12）。

我们使用了像 Nvidia A100 这样的现代 GPU，目前这类硬件尚未普及。尽管存在计算资源上的限制，我们仍相信研究人员和实践者可以从我们建立的指导原则中受益，从而更好地利用他们现有的资源。

对于合成目标分布，我们最多只测试到了一百维 。虽然这已经是一个相当高的维度，但我们的逐步分析需要大量的计算资源，在 Nvidia A100 GPU 上总共消耗了超过 4000 小时的计算时间。我们预计其中一些发现也适用于更高维度的情形，但将其扩展到更高维将留待未来研究。

一些经验丰富的专家可能会认为我们的一些发现是显而易见的。然而，我们认为，除非有人像我们一样进行直接而全面的研究，否则这些技术的实际效果仍然是假设性的。

E Real-NVP 架构细节（Details of Real-NVP Architecture）

我们使用了具有仿射耦合层 （affine coupling layers）的 Real-NVP [27] 流模型。我们定义每个耦合层由两个变换组成，其中一次变换对应于对一部分潜在变量进行仿射变换。

例如，如果第 k 层的输入变量为 z(k)，则第一个变换定义如下：

K.1 非合成模型的参考样本（Reference Samples for Non-Synthetic Models）

对于非合成模型，我们需要参考样本以评估边缘-Wasserstein 距离。为此，我们并行运行了 NumPyro 中带预处理的 NUTS （根据初步实验表现最佳）进行采样，共运行 214 条链，每条链运行 10K 次迭代，并进行抽稀处理，最终总共收集一百万个样本。

然而，由于我们选择的模型具有较为复杂的几何结构（见附录 L.2），我们发现其中若干链陷入了低概率区域（表现为在整个运行过程中样本重复），未能收集到有效样本。这种情况尤其出现在以下三个模型中：

• Exoplanet 模型
• Pharmacokinetics 模型
• Sparse Linear Regression 模型

（我们使用的是 32 位浮点数精度。在使用 64 位的情况下，性能可能会有所改善。）

在生成参考样本时，我们剔除了所有陷入低概率区域的链。这样每个模型至少保留了 497,395 个有效样本。因此，我们统一使用 497,395 个样本作为所有非合成模型的参考样本。

在进行评估时，我们尽可能多地使用方法生成的样本。如果某个方法生成的样本超过 497,395 个，则从中取前 497,395 个；若少于该数量，则全部使用。

L 目标分布细节（Target Details） L.1 合成目标分布（Synthetic Targets）

病态高斯分布（Ill-conditioned Gaussian）
这是一个零均值的多维高斯分布，其协方差矩阵的特征值从伽马分布（shape = 0.5, scale = 1）中采样，并通过一个随机正交矩阵进行旋转，从而引入高度相关性。我们使用 Inference Gym 的实现来确保目标分布可复现 [88]。

香蕉分布（Banana）[39]
该分布在二维空间中对标准正态分布进行了非线性变换：对第二维应用一个非线性变换，使其依赖于第一维的值。该变换包含一个曲率参数，它根据第一维的值调整第二维的尺度。当指定维度大于 2 时，其余维度设为标准正态变量。

对于 d 维情况，该分布定义如下：

L.2 非合成模型（Non-synthetic Models）

我们还在一些目标密度无法用闭合形式表示、也无法获得真实样本的模型上比较了 flow VI 与 HMC 方法的表现。这些模型的具体细节如下：

Exoplanet 模型（维度 = 7）
该模型描述了一颗系外行星从其宿主恒星前方经过（即发生凌星现象）时，恒星亮度下降的过程。通过分析恒星的光变曲线（light curve），该模型可以推断出行星的属性，例如其大小和轨道周期。该模型包括描述凌星形状的函数及其相关的行星参数。

控制凌星现象和系外行星特性的潜在变量具有以下先验分布：

原文链接： https://arxiv.org/pdf/2412.08824?