[论文解读] Understanding and Mitigating Exploding Inverses in Invertible Neural Networks

Jens Behrmann, Paul Vicol+3 more|arXiv (Cornell University)|Jun 16, 2020

Model Reduction and Neural Networks参考文献 59被引用 18

一句话总结

本文识别并解决了可逆神经网络（INNs）中反函数爆炸的问题，该问题导致数值非可逆性，从而破坏了归一化流和内存高效反向传播等关键应用。作者推导了INN构建模块的双-Lipschitz性质，并提出一种正则化方案以稳定局部可逆性，以及一种Lipschitz约束架构以实现全局可逆性，显著提升了生成和判别任务中的稳定性。

ABSTRACT

Invertible neural networks (INNs) have been used to design generative models, implement memory-saving gradient computation, and solve inverse problems. In this work, we show that commonly-used INN architectures suffer from exploding inverses and are thus prone to becoming numerically non-invertible. Across a wide range of INN use-cases, we reveal failures including the non-applicability of the change-of-variables formula on in- and out-of-distribution (OOD) data, incorrect gradients for memory-saving backprop, and the inability to sample from normalizing flow models. We further derive bi-Lipschitz properties of atomic building blocks of common architectures. These insights into the stability of INNs then provide ways forward to remedy these failures. For tasks where local invertibility is sufficient, like memory-saving backprop, we propose a flexible and efficient regularizer. For problems where global invertibility is necessary, such as applying normalizing flows on OOD data, we show the importance of designing stable INN building blocks.

研究动机与目标

• 调查可逆神经网络（INNs）的数值不稳定性，特别是其在实践中破坏可逆性的反函数爆炸现象。
• 通过研究常见INN构建模块（如耦合层）的双-Lipschitz性质，分析反函数爆炸的根本原因。
• 解决需要局部可逆性的应用中的不稳定性（如内存高效反向传播），以及需要全局可逆性的应用（如分布外数据上的归一化流）。
• 开发实用且有效的正则化与架构约束，确保在各类INN应用场景中实现稳定的反函数计算。

提出的方法

• 推导标准INN组件（如仿射和加法耦合层）的理论双-Lipschitz边界，以理解其稳定性特性。
• 提出一种灵活且可微的正则化器，用于在训练过程中强制实现局部可逆性，尤其适用于内存高效反向传播。
• 引入Lipschitz约束架构作为实现全局可逆性的解决方案，确保反函数雅可比奇异值有界。
• 在CIFAR-10、CelebA和Flow-GAN基准上实证验证正则化与架构约束的有效性。
• 使用标准梯度与内存节省梯度之间的夹角度量，量化训练过程中的反函数稳定性。
• 在INNs中同时应用最大似然估计（MLE）与GAN目标，并通过对比似然性（BPD）与样本质量（FID、Inception Score）来评估稳定性与性能。

实验结果

研究问题

• RQ1 为何常用INN架构在理论上可逆，却在实践中无法保持数值可逆性？
• RQ2 INN构建模块的双-Lipschitz性质如何影响反函数映射的稳定性？
• RQ3 正则化是否能有效稳定仅需局部可逆性的任务中的反函数映射，如内存高效反向传播？
• RQ4 为确保全局可逆性（尤其在归一化流模型的分布外数据上），需要哪些架构约束？
• RQ5 训练目标的选择（如MLE与GAN）如何影响INN反函数的数值稳定性及下游性能？

主要发现

• 常用INN架构（包括仿射耦合流）存在反函数爆炸问题，导致在分布内与分布外数据上均出现重建误差、NaN值，以及变量变换公式的失效。
• 使用标准反向传播训练的仿射耦合模型，其标准梯度与内存节省梯度之间的夹角迅速增大，导致数值非可逆性（夹角 → NaN），而加法模型则保持稳定。
• 通过有限差分或归一化流目标进行正则化，可使梯度夹角保持较小，从而实现稳定的内存高效训练。
• 仅使用GAN目标训练INNs会导致反函数不稳定，即使重建误差较小，实际中仍观察到无限比特每维（BPD），表明似然估计失效。
• Flow-GAN目标（结合MLE与GAN损失）相比纯GAN训练，实现了更稳定的反函数（BPD ≈ 4.21）与更优的样本质量（FID = 420），而纯GAN训练的FID为850，BPD为∞。
• 即使条件数稳定且无重建误差，仅使用GAN目标训练的模型仍赋予极高的BPD值，表明似然估计存在显著不稳定性。