https://github.com/vislearn/FFF
Lifting Architectural Constraints of Injective Flows v4 2024.04
摘要
归一化流（Normalizing Flows）显式地最大化训练数据上的全维似然。然而，真实数据通常只在较低维的流形上得到支持，导致模型在模拟噪声上消耗大量计算资源。注入流（Injective Flows）通过联合学习流形及其上的分布来解决这个问题。到目前为止，它们受到了限制性架构和/或高计算成本的限制。我们通过一种新的高效最大似然损失估计器来解除这两个限制，该估计器与自由形式的瓶颈架构兼容。我们进一步表明，简单地同时学习数据流形及其上的分布可能导致发散的解决方案，并利用这一见解来推动一个稳定的最大似然训练目标。我们在玩具数据、表格数据和图像数据上进行了广泛的实验，展示了由此产生的模型的竞争性能。
基础工作：
1 引言
生成建模是机器学习中最重要的任务之一，在视觉（Rombach et al., 2022）、语言建模（Brown et al., 2020）、科学（Ardizzone et al., 2018; Radev et al., 2022）等领域有着广泛的应用。生成建模最有动力的方法之一是最大似然训练，由于其有利的统计特性（Hastie et al., 2009）。在连续设置中，最常见的通过归一化流（Rezende & Mohamed, 2015; Dinh et al., 2015; Kobyzev et al., 2021）实现精确的最大似然训练，归一化流参数化了一个可追踪变量变化的精确可逆函数（对数行列式项）。这通常引入了模型表达性和计算成本之间的权衡，其中训练和采样最便宜的网络，如耦合块架构，需要非常特殊构造的函数，这可能限制了表达性（Draxler et al., 2022）。此外，归一化流保留了输入的维度，需要与数据空间相同维度的潜在空间。
由于流形假设（Bengio et al., 2013），该假设认为真实数据位于高维数据空间中的低维流形上，因此只在低维流形上建模分布并将流形外的偏差视为无信息的噪声更为高效。先前的工作如Caterini et al. (2021); Brehmer & Cranmer (2020) 通过特殊构造的瓶颈架构（被称为“可逆自动编码器”（Teng & Choromanska, 2019）或“注入流”（Kothari et al., 2021））将归一化流限制在低维流形上，其中编码器和解码器共享参数。这些注入流通过某种形式的最大似然训练进行优化。这不是一个理想的设计决策，因为这些模型中使用的限制性架构最初是为了在归一化流中进行可追踪的变量变化计算而设计的，但在存在瓶颈的情况下，这种计算是不可能的（Brehmer & Cranmer, 2020）。因此，我们建议放弃限制性构造（如耦合块），转而使用无约束的编码器和解码器，并引入一种新技术来绕过计算变量变化。这极大地简化了模型的设计，并使其更具表达性。
我们建立在矩形流（Caterini et al., 2021）使用的无偏最大似然估计器上，以近似变量变化项的梯度。我们通过用一个高效的单步估计器替换迭代共轭梯度，大大简化了估计器。这很快：一个批次可以在大约两倍于（或更少）仅训练重建损失的自动编码器的时间处理。此外，我们对注入流做出了一个新颖的观察：由于解码函数中发散曲率的可能性，简单地使用最大似然训练是不明确的。为了解决这个问题，我们提出了对最大似然估计器的修改，以抵消发散曲率的可能性。我们称我们的模型为自由形式注入流（FIF）。free-form injective flow (FFIFs).
总结一下，我们做出了以下贡献：
- 我们引入了一个高效的最大似然估计器，用于自由形式的注入流，并首次使用它来训练一个无约束的注入流（第4.1节）。
- 我们识别了在瓶颈存在时简单应用最大似然训练的病态行为，并提供了一个解决方案，以避免这种行为同时保持计算效率（第4.2节）。
- 我们在玩具数据、表格数据和图像数据上超越了之前的注入流，并展示了与生成自动编码器竞争的性能（第5节）。我们提供了实现我们模型和重现我们结果的代码，网址为https://github.com/vislearn/FFF。
2 相关工作
注入流联合学习一个流形并在该流形上最大化似然。后者需要估计变换的雅可比行列式来计算变量变化。传统上，这个行列式的有效计算对归一化流架构施加了两个主要限制：首先，潜在空间必须与数据空间在维度上匹配，排除了瓶颈架构。其次，归一化流被限制在某些功能形式上，如耦合和自回归块。下面我们概述了克服这些问题的现有方法以及我们的解决方案如何比较。
低维潜在空间 一组方法试图使用全维归一化流，并通过一些额外的正则化或架构约束，使得潜在空间的一个子空间对应于流形。一种策略是在数据中添加噪声以使其成为全维分布，然后去噪到流形上（Horvat & Pfister, 2021; Loaiza-Ganem et al., 2022）。另一种方法限制非流形潜在维度的方差较小（Beitler et al., 2021; Silvestri et al., 2023; Zhang et al., 2023）。
其他方法通过将训练变成两步程序来规避问题。首先，训练一个自动编码器以重建损失，然后训练一个归一化流以学习得到的潜在分布。在这一系列工作中，Brehmer & Cranmer (2020) 和 Kothari et al. (2021) 使用注入流，而 B¨ohm & Seljak (2022) 使用无约束网络作为自动编码器。Ghosh et al. (2020) 进一步正则化解码器。
保形嵌入流（Ross & Cresswell, 2021）通过进一步限制架构仅限于保形变换，确保行列式分解为每个块的贡献。Cramer et al. (2022) 使用等距自动编码器，使得变量变化变得微不足道。然而，由此产生的变换非常受限，不能表示任意流形。
与我们最相似的工作是矩形流（Caterini et al., 2021），它通过一个迭代的无偏估计器估计对数行列式的梯度。由此产生的方法训练起来相当慢，并且使用限制性的注入流。
无约束的归一化流架构 几项工作试图减少典型归一化流架构施加的约束，允许使用自由形式的网络。然而，所有这些方法仅适用于全维架构。FFJORD（Grathwohl et al., 2019）是一种连续归一化流（Chen et al., 2018b），它随机估计变量变化。残差流（Behrmann et al., 2019; Chen et al., 2019）使残差网络可逆，但需要昂贵的迭代估计器通过最大似然进行训练。自归一化流（Keller et al., 2021）和相对梯度优化（Gresele et al., 2020）估计神经网络中使用的矩阵的最大似然梯度，但限制架构仅使用没有跳跃连接的方形权重矩阵。在平行工作中，我们提出了基于自由形式神经网络的全维归一化流的扩展，称为自由形式流（Draxler et al., 2024）。
近似最大似然 许多方法优化了全维最大似然的一些界限，特别是变分自动编码器（Kingma & Welling, 2014）及其变体。Cunningham et al. (2020) 也优化了似然的一个变分下界。其他方法通过联合优化重建损失和流形上最大似然的某种近似来适应注入流框架：Kumar et al. (2020); Zhang et al. (2020) 通过其Frobenius范数近似雅可比的对数行列式。熵AE（Ghose et al., 2020）通过最近邻估计器最大化潜在分布的熵，同时限制其方差，导致高斯潜在空间。此外，还有其他正则化自动编码器潜在空间的方法，这些方法不是基于最大似然的，例如对抗性方法（Makhzani et al., 2015）。
与上述方法相反，我们的方法联合学习流形并在其上最大化似然，具有无约束的架构，可以轻松容纳低维潜在空间。
3背景
在维度之间的变量变换 最大似然目标是由变量变换定理导致的，用于训练归一化流模型，仅在映射到相等维度空间时才有明确定义。微分几何中的一个结果（Krantz＆Parks，2008）使我们能够通过以下公式将变量变换定理推广到非等维度变换：
在这个公式中，f和g是一致的，并且上标表示导数：g′(f(x))是在f(x)处评估的g的雅可比矩阵。请注意，由于pX是通过g将潜在分布pZ推前得到的，因此该公式仅对位于解码器流形上的x有效（更多细节请参见附录A）。与全维度情况不同，目前没有已知的架构能够同时表示任意流形并允许对等式(2)进行高效精确计算（参见第2节）。
矩形流 最小化等式(2)的负对数，并添加拉格朗日乘子以限制数据点与解码器流形之间的距离，得到以下每样本损失项：
该计算替代了损失函数中的对数行列式项。尽管它的值不同，但其梯度与原损失函数（公式(3)）的梯度相同。接下来，我们基于编码器是解码器的（近似）逆这一见解，提出了该估计器的一个显著改进版本。
4 自由形式注入流 (FIF)
我们对矩形流的修改有三个方面：首先，我们使用一个无约束的自动编码器架构（没有严格参数化的可逆函数）；其次，我们引入了一个更高效的替代估计器；第三，我们修改了替代估计器，以避免与高曲率流形相关的病态行为。我们的每个样本损失函数如下：
其中，z = f(x)。请注意，替代项前面的负号来自将对数行列式的梯度发送到编码器而不是解码器。我们将在4.1和4.2节推导和解释损失的这种表述。
4.1 简化替代估计器
我们通过在公式（2）中引入一个新的对数行列式项的替代估计器，大大简化了矩形流的优化过程，该替代估计器使用编码器的雅可比矩阵来近似解码器的逆雅可比矩阵。这使得替代估计器可以在一次传递中计算，避免了昂贵的共轭梯度迭代。
我们通过在公式（4）中展开导数来实现这一点：
每个求和项只需从自动微分中获得的两个向量-雅可比/雅可比-向量积即可计算。这是对原始矩形流的迭代共轭梯度方法的显着改进，原方法需要最多2(d+1)次这样的积来确保收敛（Caterini等人，2021年）。与仅重构损失相比，我们测量墙上的时间约为1.5倍至2倍，与潜在维度无关。请注意，替代估计器只在f和g相对于重构损失（至少近似）最优时才准确。我们观察到实践中的稳定训练，验证了这一假设。
4.2 在瓶颈存在的情况下最大似然估计的问题
矩形流的训练结合了重构项和似然项。我们可能会问，如果我们只用似然项进行训练会发生什么，类比于归一化流。在这种情况下，我们的损失将是:
因此，损失的下限受投影到流形上的数据熵的限制：
与标准的规范化流优化不同，在这里右边的项是固定的，熵取决于模型学习的投影。因此，模型可以修改投影，使熵尽可能低。我们将这种病态情况分为两种情况：
1. 模型流形与数据流形不对齐而是与其相交。例如，Brehmer和Cranmer（2020）讨论了一个线性模型学习将数据分布投影到流形上的单个点的情况，从而将其熵减少到无穷大，这是最低可能的值。据我们所知，可以通过添加噪声和一个具有足够高权重的重构损失来解决这个问题。在附录C中，我们对线性模型进行了证明，并描述了其解决方案，如果，其中是数据协方差矩阵的最小特征值，则解决方案与PCA相同。
2. 模型流形通过使用高曲率来集中数据，见图2（左）。这种新发现的病态情况仅发生在非线性模型中，因此Brehmer和Cranmer（2020）在他们的线性示例中没有注意到这一效果。重要的是，这不能通过添加重构损失来解决。
大多数现有的注入流通过两阶段训练来避免这种情况，首先学习一个投影，然后在潜在空间中学习投影数据的分布。为了能够共同学习流形及其上的最大似然密度，我们需要找到一种解决病态的方法。
5 实验
在本节中，我们测试了所提出的模型的经验性能。首先，我们展示了我们的模型在表格数据上比矩形流快得多。其次，我们展示了它在生成图像方面优于先前的SOTA注入流。最后，我们在Pythae图像生成基准Chadebec等人（2022年）上与其他生成自动编码器进行比较，在某些类别中获得了最佳的FID分数。
实现细节：在实现迹估计器时，我们必须做出一些选择。简而言之，i）我们选择将对数行列式梯度的公式化为编码器而不是解码器，因为在实践中更稳定；ii）我们按照f'（x）g'(z)的顺序执行迹，因为这样可以减少方差（由于迹的循环性质，两种排序都是有效的，但由于f'(x)g'(z)是d×d矩阵，而g'(z)f'(x)是D×D矩阵，前者通常更容易估计）；iii）我们使用正向和反向模式自动微分的混合，因为这与我们的估计器兼容；iv）我们在迹估计器中使用了正交高斯噪声，以减少方差。有关这些选择的全面理由，请参见附录D。
5.1 表格数据
我们在Papamakarios等人（2017年）使用的四个表格数据集上评估了我们的方法，使用相同的数据拆分，并与已发布的矩形流结果（Caterini等人，2021年）进行比较，见表1。我们采用了该工作中的“类似FID指标”，该指标计算了具有相等均值和协方差的高斯分布与测试数据和模型生成的数据之间的Wasserstein-2距离。这是生成数据集和测试数据集的均值和协方差矩阵之间差异的度量。我们在所有数据集上都优于矩形流，除了GAS。此外，我们在训练时间上也看到了与FIF和矩形流（使用已发布的代码）在相同硬件上重新运行之间的加速，加速比为1.5至6倍。完整的实验细节见附录E.2。
我们还在附录中的表6和表7中进行了消融研究，解析了我们方法的三个单独组成部分的效果：公式（10）的替代项，公式（15）中高曲率解的修正，以及自由形式架构的使用。与表1相比，我们发现使用在流形外部编码器雅可比矩阵来估计负对数似然是自由形式架构性能良好的关键，因为在流形内部变体会发散（见第4.2节）。这对基于耦合流的注入架构并非如此，表明了通过架构的稳定正则化。
5.2 与注入流的比较
我们在CelebA图像（Liu等人，2015年）上将FIF与先前的注入流进行了比较，见表2。我们的模型在生成图像的质量上显着提高，以Frechet Inception Distance（FID）（Heusel等人，2017年）和Inception Score（IS）（Salimans等人，2016年）为指标。前者将生成样本与一组参考样本进行比较，通过计算两个高斯分布拟合到各自样本集的某种嵌入之间的Wasserstein-2距离来完成。后者通过生成样本中类标签分布的熵来衡量多样性，其中类标签由某个预训练分类器提供。来自该模型的样本显示在图1中。
为了公平比较，我们使用作者提供的代码在相同的硬件上以相等的墙上时间训练每个模型。先前工作的架构主要受到大多数层需要是可逆的并且具有可处理的雅可比行列式的约束的影响。我们的损失公式（7）不对架构施加这些约束，因此我们可以使用一个现成的卷积自动编码器，并在潜在空间中添加额外的全连接层。详情请参见附录E.3。
5.3 与生成自动编码器的比较
由于自由形式注入流（FIF）不需要任何特定的架构，我们将比较扩展到更广泛的生成自动编码器范围。这是一种将训练数据编码为标准正态分布的通用瓶颈架构，因此在训练后，解码器可以用作生成器。
最近，Chadebec等人（2022年）提出了Pythae基准，用于比较图像生成中的生成自动编码器。他们使用两种不同的架构评估不同的训练方法，分别在MNIST（LeCun等人，2010年）（数据D = 784，潜在d = 16）、CIFAR10（Krizhevsky，2009年）（D = 3072，d = 256）和CelebA（Liu等人，2015年）（D = 12288，d = 64）上。所有模型都在相同的有限计算预算下进行训练。该基准的目标是提供对不同模型的公平比较，而不是实现SOTA图像生成结果，因为这将需要更多的计算资源。
正如表3所示，我们的模型在该基准上表现强劲，在CelebA上以标准正态分布中的潜在代码为样本，并在使用训练数据拟合的高斯混合模型进行采样时，在ResNet架构中的Fréchet Inception Distance（FID）（Heusel等人，2017年）上达到了SOTA，而在两种架构中采样时也是如此。同时，Inception Score（IS）（Salimans等人，2016年）很高，表明了很高的多样性。在我们的模型没有超越竞争对手的组合中，FIF仍然实现了可比的FID和高的Inception Score。FIF在其他数据集上的表现也很强劲，请参见附录E.4。
在基准中，对于每种方法，训练十种超参数配置，并报告FID最佳的模型。对于我们的方法，我们选择变化β = 5、10、15、20、25和Hutchinson采样数量K = 1、2。我们发现性能对这些选择具有鲁棒性，并在附录E.4中提供了训练过程的所有细节。
结论
本文提出了一种计算效率高的解决方案，可以共同学习流形及其上的分布，我们将其称为自由形式注入流（FIF）。我们i）显著改进了一个现有的跨维度变量变换梯度估计器，ii）指出它可以应用于无约束的架构，iii）分析了联合流形和最大似然训练的问题并提供了解决方案，以及iv）在玩具、表格和图像数据集上实现和测试了我们的模型。我们发现该模型实用且可扩展，在表现上优于可比的注入流，并在Pythae基准上显示出与其他自动编码器生成模型类似或更好的性能。
未来的研究仍然存在一些理论和实践问题：我们发现了一个先前忽视的共同学习流形和最大似然的问题。我们在第4.2节提出了一个解决方案，可以提供高质量的模型，但需要进一步调查以深入理解。在训练后将GMM拟合到潜在空间可以提高图像数据的性能，这表明我们的潜在分布并不完全是高斯的。我们通常发现，在潜在空间中具有更多全连接层的架构具有更高的高斯潜在分布，这表明更大的模型更少受到这个问题的影响。我们将潜在的理论或实践改进留给未来的工作。
热门跟贴