Density estimation using real nvp

发布时间 2023-04-12 10:01:14作者: silly丶

Density estimation using real nvp

参考资料:

  • 苏剑林. (Aug. 26, 2018). 《细水长flow之RealNVP与Glow:流模型的传承与升华 》[Blog post]. Retrieved from https://kexue.fm/archives/5807
  • 《Density estimation using real nvp》

简介

生成概率模型不仅具有创建新内容的能力,而且还具有广泛的重建相关应用,包括修补、去噪、着色和超分辨率。

由于感兴趣的数据通常是高维和高度结构化的,这一领域的挑战是建立足够强大的模型来捕捉其复杂性,同时仍然是可训练的。我们通过引入 real-valued non-volume preserving(real NVP,实值非体积保持)转换 来解决这一挑战,这是一种易于处理且具有表现力的高维数据建模方法。

该模型能够对数据点进行高效准确的推断、采样和对数密度估计。此外,本文提出的体系结构可以根据该模型提取的层次特征精确有效地重建输入图像。

相关工作

variational autoencoder algorithm:变分自编码器算法通过利用重参数化技巧,同时学习一个将高斯隐变量 \(z\) 映射到样本 \(x\) 的生成网络,以及一个匹配的将样本 \(x\) 映射到语义上有意义的潜在表示 \(z\) 的近似推理网络。它成功地利用了深度神经网络中反向传播的最新进展,这使它被应用于语音合成到语言建模等多个应用。尽管如此,推理过程中的近似限制了它学习高维深度表示的能力,这激励了最近改进近似推理的工作。

Generative Adversarial Networks:生成对抗网络(GANs)可以通过完全避免最大似然原则来训练任何可微生成网络。生成网络与鉴别器网络相关联,其任务是区分样本和真实数据。这个鉴别器网络以对抗的方式提供训练信号,而不是使用难以处理的对数似然。成功训练的GAN模型可以持续生成清晰且逼真的样本。然而,测量生成样本中的多样性的指标目前是难以解决的。此外,训练过程的不稳定性需要仔细的超参数调优,以避免发散行为。

训练这样一个将隐变量 \(z \sim p_Z\) 映射到样本 \(x \sim p_X\) 的生成网络 \(g\),在理论上不需要像GANs那样的鉴别器网络,也不需要像变分自编码器那样的近似推理。事实上,如果 \(g\) 是双射的,它可以使用变量变换公式通过最大似然来训练:

\[p_X(x)=p_Z(z)\left| \text{det}\left( \frac{\partial g(z)}{\partial z^T} \right) \right|^{-1} \tag{1} \]

模型

Change of variable formula

给定一个观测数据变量 \(x \in X\)。一个隐变量 \(z \in Z\) 上的简单先验概率分布 \(p_Z\),一个双射 \(f:X\rightarrow Z(g = f^{-1})\),变量变换公式定义了 \(x\) 上的模型分布:

\[p_X(x)=p_Z(f(x))\left| \text{det} \left( \frac{\partial f(x)}{\partial x^T} \right) \right| \tag{2} \]

\[log(p_X(x))=log(P_Z(f(x)))+log\left( \left| \text{det}\left( \frac{\partial f(x)}{\partial x^T} \right) \right| \right) \tag{3} \]

其中 \(\frac{\partial f(x)}{\partial x^T}\) 是函数 \(f\) 关于 \(x\) 的的雅克比矩阵

在隐空间中绘制了一个样本 \(z \sim p_Z\),它的逆的像 \(x=f^{-1}(z)=g(z)\) 生成了一个原空间的样本。计算点 \(x\) 上的密度,可以通过计算它的像 \(f(x)\) 密度并乘以相关的雅克比行列式 $ \text{det} \left( \frac{\partial f(x)}{\partial x^T} \right)$ 得到。如图1所示

image-20230411193500228

耦合层

通过对函数 \(f\) 的精心设计,可以学习到一个双射模型,该模型既易于处理,又非常灵活。由于计算变换的雅可比行列式对于有效地训练使用这个原理是至关重要的,这项工作利用了一个简单的观察,即三角形矩阵的行列式可以有效地计算为其对角线项的乘积。

我们将通过叠加一系列简单的双射函数来构建一个灵活且易于处理的双射函数。在每个简单的双射中,输入向量的\(x_{d+1:D}\)部分使用一个函数更新,这个函数很容易求逆,但它以复杂的方式依赖于输入向量的剩余部分。我们把这些简单的双射称为一个仿射耦合层。给定 \(D\) 维输入\(x\)\(d<D\),仿射耦合层的输出 \(y\) 符合方程:

\[\begin{align} y_{1:d}&=x_{1:d} \tag{4}\\ y_{d+1:D}&=x_{d+1:D}\odot exp(s(x_{1:d}))+t(x_{1:d}) \tag{5} \end{align} \]

其中 \(s\)\(t\) 表示缩放和平移,两个函数均为 \(R^d \mapsto R^{D-d}\)\(\odot\) 代表哈达玛积乘机。

具体如图2(a)所示, \(x_1\) 对应 \(x_{1:d}\)\(x_2\) 对应 \(x_{d+1:D}\)。公式(5)相当于对 \(x_2\) 进行一个仿射变换。

image-20230411194100139

雅克比矩阵

该变换的雅克比矩阵如下:

\[\frac{\partial y}{\partial x^T}=\begin{bmatrix} \mathbb{I}_d & 0\\ \frac{\partial y_{d+1:D}}{\partial x^T_{1:d}} & diag(exp[s(x_{1:d})]) \end{bmatrix} \tag{6} \]

该雅克比矩阵的行列式为 \(exp(\sum_js(x_{1:d})_j)\)。由于计算耦合层运算的雅可比行列式不涉及计算s或t的雅可比行列式,因此这些函数可以任意复杂。在本文我们使用深度卷积网络。

计算这些耦合层的逆函数并不比前向传播复杂,这意味着对于这个模型,采样和推理一样有效。其公式表达如下:

\[\begin{cases} y_{1:d}&=x_{1:d}\\ y_{d+1:D}&=x_{d+1:D}\odot exp(s(x_{1:d}))+t(x_{1:d})\\ \end{cases}\\\tag{7} \]

\[\Leftrightarrow \begin{cases} x_{1:d}&=y_{1:d}\\ x_{d+1:D}&=(y_{d+1:D} - t(y_{1:d})) \odot exp(-s(y_{1:d}))\\ \end{cases}\\\tag{8} \]

具体如图2(b)所示。主要是 \(y_1=x_1\),导致倒推十分方便。

带掩蔽的卷积

flow中的两个操作:

  1. 将输入分割为两部分 \(x_1,x_2\) ,然后输入到耦合层中,而模型中的 \(s,t\) 事实上只对 \(x_1\) 进行处理;
  2. 特征输入耦合层之前,要随机打乱原来特征的各个维度(相当于乱序的特征)。这两个操作都会破坏局部相关性,比如分割操作有可能割裂原来相邻的像素,随机打乱也可能将原来相邻的两个像素分割得很远。

RealNVP约定分割和打乱操作,都只对“通道”轴执行。也就是说,沿着通道将输入分割为 \(x_1,x_2\)后,\(x_1\)还是具有局部相关性的。还有沿着通道按着同一方式打乱整体后,空间部分的相关性依然得到保留,因此在模型 \(s,t\) 中就可以使用卷积了。

沿着通道轴进行分割,不损失空间上的局部相关性

沿着通道轴进行分割,不损失空间上的局部相关性

在RealNVP中,将输入分割为两部分的操作称为mask,因为这等价于用0/1来区别标注原始输入。

除了前面说的通过通道轴对半分的mask外,RealNVP事实上还引入了一种空间轴上的交错mask,如下图。这种mask称为棋盘式mask(格式像国际象棋的棋盘)。这种特殊的分割也保留了空间局部相关性,原论文中是两种mask方式交替使用的,但这种棋盘式mask相对复杂,也没有什么特别明显的提升,所以在Glow中已经被抛弃。

沿着空间轴交错(棋盘)分割,也是一种保持空间局部相关性的方案

沿着空间轴交错(棋盘)分割,也是一种保持空间局部相关性的方案

分区可以使用二进制掩码 \(b\) 实现,使用 \(y\) 的函数形式如下:

\[y=b\odot x + (1-b) \odot (x \odot \exp(s(b\odot x))+t(b \odot x)) \tag{9} \]

耦合层组合

尽管耦合层功能强大,但它们的正向转换会使某些组件保持不变,如下图所示。

简单的耦合使得其中一部分仍然保持恒等,信息没有充分混合

简单的耦合使得其中一部分仍然保持恒等,信息没有充分混合

这个困难可以通过交替组合耦合层来克服,这样在一个耦合层中保持不变的组件在下一个耦合层中更新。如下图所示。图4(a)是原论文中的图。此处觉得下图更加清晰,所以也贴过来。

通过交叉耦合,充分混合信息,达到更强的非线性

image-20230412091022180

最后的得到的 雅克比行列式也是容易计算的

\[\frac{\partial (f_b \circ f_a)}{\partial x_a^T}(x_a)=\frac{\partial f_a}{\partial x_a^T}(x_a)\cdot\frac{\partial f_b}{\partial x_b^T}(x_b=f_a(x_a)) \tag{10} \]

\[\det(A \cdot B)=\det(A)\det(B) \tag{11} \]

类似的,它的逆函数也能容易计算得到

\[(f_b \circ f_a)^{-1}=f_a^{-1} \circ f_b^{-1} \tag{12} \]

多尺度结构

squeeze操作

RealNVP使用squeeze操作实现了一个多尺度架构:对于每个通道,它将图像划分为 \(2 \times 2 \times c\) 形状的子正方形,然后将它们重新塑造为 \(1 \times 1 \times 4c\) 形状的子正方形。压缩操作将 \(h \times w \times c\) 张量转换为 \(\frac{h}{2} × \frac{w}{2} × 4c\) 的张量,有效地用空间大小交换通道数量。

squeeze操作图示,其中2x2的小区域可以换为自定义大小的区域

squeeze操作图示,其中2x2的小区域可以换为自定义大小的区域

image-20230412092458218

多尺度结构

在每个尺度上,我们将几个操作组合成一个序列:我们首先应用三个使用棋盘掩码的耦合层,然后执行压缩操作,最后应用三个使用通道掩码的耦合层。选择通道级掩码,这样得到的分区就不会与之前的棋盘掩码冗余(参见图3)。对于最后一个尺度,我们只应用四个交替棋盘掩码的耦合层。

在所有耦合层中传播 \(D\) 维向量会很麻烦,这涉及到计算和内存成本,以及需要训练的参数数量。RealNVP的每一步的多尺度操作直接将数据尺寸减少到原来的一半。如图3(b)所示。

原始输入经过第一步flow运算(“flow运算”指的是多个仿射耦合层的复合)后,输出跟输入的大小一样,这时候将输入对半分开两半\(z_1,z_2\) (自然也是沿着通道轴),其中 \(z_1\) 直接输出,而只将 \(z_2\) 送入到下一步flow运算,后面的依此类推。比如图中的特例,最终的输出由 \(z_1,z_3,z_5\) 组成,总大小跟输入一样。

RealNVP中的多尺度结构图示

RealNVP中的多尺度结构图示

具体公式如下:

\[h^{(0)}=x\\ (z^{(i+1)},h^{(i+1)})=f^{(i+1)}(h^{(i)}) \]

最终输出的先验分布