Normalizing Flow

flow的核心思想就是这个分布变换的公式，如果$y=f(x)$，且$f$是可逆的，则

$$\begin{gathered} p_x(x)=p_y(f(x))\ast |\det Jf(x)| \\ {p}_y(y)=p_x(f^{-1}(y))\ast |\det J{f}^{-1}(y)| \end{gathered}$$

那如果有很多个$f$，那么取log之后我们只需简单地加起来就好了：

$$\begin{array}{rl}{\mathbf{z}}_K& =f_K\circ \dots \circ f_2\circ f_1({\mathbf{z}}_0)\\ \log q_K({\mathbf{z}}_K)& =\log q_0({\mathbf{z}}_0)-\sum _{k=1}^K\log \det \left|\frac{\partial f_k}{\partial {\mathbf{z}}_k}\right|\end{array}$$

想要对这个分布变换了解更多的可以看我前一篇文章：
理解Jacobian矩阵与分布变换

NICE: Additive coupling layer

从分布变换的公式可以看出，想要处理好这种变换，首先要保证f是可逆的，其次，f的Jacobian的行列式也要好求才行，而最好求的行列式自然就是三角行列式，等于对角线的。

那么现在介绍最基本的做法就是NICE的做法：首先对于D维数据的x，将其随意地划分为两部分，$x_1,x_2$,并做变换：

$$\begin{array}{rl}& {\mathbf{y}}_1={\mathbf{x}}_1\\ & {\mathbf{y}}_2={\mathbf{x}}_2+\mathbf{m}({\mathbf{x}}_1)\end{array}$$

其中m是一个任意的MLP函数，通过这样的变换，我们发现这个函数是可逆的，即：

$$\begin{array}{rl}& {\mathbf{x}}_1={\mathbf{y}}_1\\ & {\mathbf{x}}_2={\mathbf{y}}_2-\mathbf{m}({\mathbf{x}}_1)\end{array}$$

事实上随后的改进都只是进一步地将这个可逆函数变得更加地“复杂”，这样的加性确实是简单了点。

它的基本原理很简单，就是将x和y划分成两块，其中$x_1=x_{1:d}$前d个元素，$x_2=x_{d+1:D}$后面的元素。（所以这个似乎只能处理x和y都是高维的情况）。注意观察，y2和x2的关系其实是线性，而且利用了x1和y1的等价关系，所以用x2表示y2的时候，直接把x1换成y1就可以了，因此这个函数非常容易求逆。其实我觉得这个东西的本质其实应该是利用了一个Z的共享变量，来构造的可逆函数

$$\begin{array}{l}Z\\ /\backslash \\ X-Y\end{array}$$

于是

$$y_2=x_2+m(z)$$

而这个函数求导也很简单

$$\frac{\partial \mathbf{y}}{\partial \mathbf{x}}=\left[\begin{array}{cc}\frac{\partial y_1}{\partial x_1}& \frac{\partial y_1}{\partial x_2}\\ \frac{\partial y_2}{\partial x_1}& \frac{\partial y_2}{\partial x_2}\end{array}\right]=\left[\begin{array}{cc}{\mathbb{I}}_{1:d}& 0\\ \frac{\partial m(x_1)}{\partial x_1}& {\mathbb{I}}_{d:D}\end{array}\right]$$

神奇的事情出现了，这个可逆函数的Jacobian矩阵居然只是一个三角矩阵，他的行列式就等于对角线的乘积为1，而其对数为0,

Real NVP: Affine Coupling layers

只是加性就太简单了，所以我们变得复杂一点，这是Real NVP中提出的做法：

$$\begin{array}{rl}& {\mathbf{y}}_1={\mathbf{x}}_1\\ & {\mathbf{y}}_2=\mathbf{s}({\mathbf{x}}_1){\odot \mathbf{x}}_2+t({\mathbf{x}}_1)\end{array}$$

我乘一个非线性函数$\mathbf{s}({\mathbf{x}}_1)$上去，这里$\odot$是点乘，其实他们本质上还是一个线性变换而已(所以才叫affine)，s就是斜率，t是截距。既然线性变换可以，肯定也有多项式变换等等变种，事实上已经有类似的工作，比如Neural Spline Flows这种就是一个多项式的可逆函数，这里先不说这个。我们先看看这个Affine Coupling layer的jacobian是长什么样：

$$\frac{\partial \mathbf{y}}{\partial \mathbf{x}}=\left[\begin{array}{cc}\frac{\partial y_1}{\partial x_1}& \frac{\partial y_1}{\partial x_2}\\ \frac{\partial y_2}{\partial x_1}& \frac{\partial y_2}{\partial x_2}\end{array}\right]=\left[\begin{array}{cc}{\mathbb{I}}_d& 0\\ \frac{\partial \mathbf{s}}{\partial {\mathbf{x}}_1}\otimes {\mathbf{x}}_2+\frac{\partial t}{\partial {\mathbf{x}}_1}& \mathrm{diag}(s)\end{array}\right]$$

其中之所以是对角矩阵是因为

$$\left(\begin{array}{c}{x}_1\\ ⋮\\ x_n\end{array}\right)\odot \left(\begin{array}{c}{y}_1\\ ⋮\\ y_n\end{array}\right)=\left(\begin{array}{ccc}{x}_1& \cdots & 0\\ ⋮& \ddots & ⋮\\ 0& \cdots & x_n\end{array}\right)\left(\begin{array}{c}{y}_1\\ ⋮\\ y_n\end{array}\right)$$

所以

$$\frac{\partial y_2}{\partial x_2}=\frac{\partial \mathrm{d}\mathrm{i}\mathrm{a}\mathrm{g}(\mathbf{s})x_2}{\partial x_2}=\mathrm{diag}(\mathbf{s})$$

这里一般会约束s大于0，所以一般神经网络输出的是log(s)，然后取指数变回来。

Glow 一种可逆的1x1卷积

我们刚才随机划分的方法感觉处理图片的时候怪怪的，而且就算是随机交换channel也感觉不太对，有没有更优雅的方法？Glow给出了解决的方法，我们可以引入1x1可逆卷积核来代替这个划分的操作。其实1x1卷积核本身就有随机置换的味道在里面，只不过这篇文章的贡献在于用了一个trick保证了他的可逆性。

可以看个小例子（引用来自苏剑林的博客：https://kexue.fm/archives/5807），随机置换的操作其实就是一个简单的线性变换：

$$\left(\begin{array}{c}2\\ 1\\ 4\\ 3\end{array}\right)=\left(\begin{array}{cccc}0& 1& 0& 0\\ 1& 0& 0& 0\\ 0& 0& 0& 1\\ 0& 0& 1& 0\end{array}\right)\left(\begin{array}{c}1\\ 2\\ 3\\ 4\end{array}\right)$$

我们知道一个卷积核其实是可以表达成一个矩阵乘积的，只要我们把所有channel展开成一条向量，就可以写出一个矩阵的计算公式：

$$Y=XW$$

比如，$x$是$h\ast w\ast c$的张量，c表示channel，对于1x1卷积来说，W就是一个$c\ast c$的矩阵，就是他们的乘积，实际上可以将x想象成$h\ast w$行$c$列的矩阵，然后乘以W。

所以接下来的问题只有一个，如何保证这个W是可逆的。为了构造一个一定可逆的矩阵W，我们可以利用LU分解。因为任意矩阵都可以表达成

$$W=PLU$$

其中P是置换矩阵， L是下三角矩阵，对角线元素全为1，U是上三角矩阵，所以为了保证矩阵W可逆，那么只要保证P, L,U满秩就可以了，又因为P,L一定是满秩的，所以只要保证U满秩即可。那么一个方便的方法就是：

$$W=PL(U+diag(s))$$

其中，U是严格上三角矩阵，其对角线为0，我们只需保证这个s不为0即可。于是最终我们的可逆卷积核其导数行列式的求解为：

$$\begin{gathered} \log \left|\det \left(\frac{d\mathrm{conv}2\mathrm{D}(\mathbf{h};\mathbf{W})}{d\mathbf{h}}\right)\right|=h\cdot w\cdot \log |\det (\mathbf{W})| \\ =h\cdot w\cdot \log |diag(s)|=h\cdot w\cdot sum(log(|s|)) \end{gathered}$$

这里有个绝对值是因为这个分布变换jacobian的行列式一定是大于0的。实际做的时候，P固定这，只要更新L，U和s的参数就好了。

此外，从实际的角度，如果加了BN后，其实相比加性耦合，仿射耦合效果的提升并不高，所以要训练大型的模型，为了节省资源，一般都只用加性耦合，比如Glow训练256x256的高清人脸生成模型，就只用到了加性耦合。

参考资料

Dinh L, Krueger D, Bengio Y. NICE: Non-linear Independent Components Estimation[J]. 2014, 1(2): 1–13.

Dinh L, Sohl-Dickstein J, Bengio S. Density estimation using Real NVP[J]. 2016.

Kingma D P, Dhariwal P, Francisco S. Glow: Generative Flow with Invertible 1×1 Convolutions[J]. : 1–15.

https://kexue.fm/archives/5807