Normalizing Flow
flow的核心思想就是这个分布变换的公式,如果 y = f ( x ) \displaystyle y=f( x) y=f(x),且 f \displaystyle f f是可逆的,则
p x ( x ) = p y ( f ( x ) ) ∗ ∣ det J f ( x ) ∣ p y ( y ) = p x ( f − 1 ( y ) ) ∗ ∣ det J f − 1 ( y ) ∣ p_{x} (x)=p_{y} (f(x))*|\det Jf(x)|\\ p_{y} (y)=p_{x} (f^{-1} (y))*|\det Jf^{-1} (y)| px(x)=py(f(x))∗∣detJf(x)∣py(y)=px(f−1(y))∗∣detJf−1(y)∣
那如果有很多个 f \displaystyle f f,那么取log之后我们只需简单地加起来就好了:
z K = f K ∘ … ∘ f 2 ∘ f 1 ( z 0 ) log q K ( z K ) = log q 0 ( z 0 ) − ∑ k = 1 K log det ∣ ∂ f k ∂ z k ∣ \begin{aligned} \mathbf{z}_{K} & =f_{K} \circ \dotsc \circ f_{2} \circ f_{1}(\mathbf{z}_{0})\\ \log q_{K}(\mathbf{z}_{K}) & =\log q_{0}(\mathbf{z}_{0}) -\sum ^{K}_{k=1}\log\operatorname{det}\left| \frac{\partial f_{k}}{\partial \mathbf{z}_{k}}\right| \end{aligned} zKlogqK(zK)=fK∘…∘f2∘f1(z0)=logq0(z0)−k=1∑Klogdet∣∣∣∣∂zk∂fk∣∣∣∣
想要对这个分布变换了解更多的可以看我前一篇文章:
理解Jacobian矩阵与分布变换
NICE: Additive coupling layer
从分布变换的公式可以看出,想要处理好这种变换,首先要保证f是可逆的,其次,f的Jacobian的行列式也要好求才行,而最好求的行列式自然就是三角行列式,等于对角线的。
那么现在介绍最基本的做法就是NICE的做法:首先对于D维数据的x,将其随意地划分为两部分, x 1 , x 2 \displaystyle x_{1} ,x_{2} x1,x2,并做变换:
y 1 = x 1 y 2 = x 2 + m ( x 1 ) \begin{aligned} & \mathbf{y}_{1} =\boldsymbol{x}_{1}\\ & \mathbf{y}_{2} =\boldsymbol{x}_{2} +\boldsymbol{m} (\boldsymbol{x}_{1} ) \end{aligned} y1=x1y2=x2+m(x1)
其中m是一个任意的MLP函数,通过这样的变换,我们发现这个函数是可逆的,即:
x 1 = y 1 x 2 = y 2 − m ( x 1 ) \begin{aligned} & \boldsymbol{x}_{1} =\mathbf{y}_{1}\\ & \boldsymbol{x}_{2} =\mathbf{y}_{2} -\boldsymbol{m} (\boldsymbol{x}_{1} ) \end{aligned} x1=y1x2=y2−m(x1)
事实上随后的改进都只是进一步地将这个可逆函数变得更加地“复杂”,这样的加性确实是简单了点。
它的基本原理很简单,就是将x和y划分成两块,其中 x 1 = x 1 : d \displaystyle x_{1} =x_{1:d} x1=x1:d前d个元素, x 2 = x d + 1 : D \displaystyle x_{2} =x_{d+1:D} x2=xd+1:D后面的元素。(所以这个似乎只能处理x和y都是高维的情况)。注意观察,y2和x2的关系其实是线性,而且利用了x1和y1的等价关系,所以用x2表示y2的时候,直接把x1换成y1就可以了,因此这个函数非常容易求逆。其实我觉得这个东西的本质其实应该是利用了一个Z的共享变量,来构造的可逆函数
Z / \ X − Y \begin{array}{l} \ \ \ \ \ \ \ Z\\ \ \ \ /\ \ \ \ \ \ \ \backslash \\ \ X\ -\ Y \end{array} Z / \ X − Y
于是
y 2 = x 2 + m ( z ) y_{2} =x_{2} +m(z) y2=x2+m(z)
而这个函数求导也很简单
∂ y ∂ x = [ ∂ y 1 ∂ x 1 ∂ y 1 ∂ x 2 ∂ y 2 ∂ x 1 ∂ y 2 ∂ x 2 ] = [ I 1 : d 0 ∂ m ( x 1 ) ∂ x 1 I d : D ] \frac{\partial \mathbf{y}}{\partial \mathbf{x}} =\left[\begin{array}{ c c } \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}{ c c } \mathbb{I}_{1:d} & 0\\ \frac{\partial m( x_{1})}{\partial x_{1}} & \mathbb{I}_{d:D} \end{array}\right] ∂x∂y=[∂x1∂y1∂x1∂y2∂x2∂y1∂x2∂y2]=[I1:d∂x1∂m(x1)0Id:D]
神奇的事情出现了,这个可逆函数的Jacobian矩阵居然只是一个三角矩阵,他的行列式就等于对角线的乘积为1,而其对数为0,
Real NVP: Affine Coupling layers
只是加性就太简单了,所以我们变得复杂一点,这是Real NVP中提出的做法:
y 1 = x 1 y 2 = s ( x 1 ) ⊙ x 2 + t ( x 1 ) \begin{aligned} & \mathbf{y}_{1} =\boldsymbol{x}_{1}\\ & \mathbf{y}_{2} =\mathbf{s}(\mathbf{x}_{1})\boldsymbol{\odot x}_{2} +t(\boldsymbol{x}_{1} ) \end{aligned} y1=x1y2=s(x1)⊙x2+t(x1)
我乘一个非线性函数 s ( x 1 ) \displaystyle \mathbf{s}(\mathbf{x}_{1}) s(x1)上去,这里 ⊙ \displaystyle \odot ⊙是点乘,其实他们本质上还是一个线性变换而已(所以才叫affine),s就是斜率,t是截距。既然线性变换可以,肯定也有多项式变换等等变种,事实上已经有类似的工作,比如Neural Spline Flows这种就是一个多项式的可逆函数,这里先不说这个。我们先看看这个Affine Coupling layer的jacobian是长什么样:
∂ y ∂ x = [ ∂ y 1 ∂ x 1 ∂ y 1 ∂ x 2 ∂ y 2 ∂ x 1 ∂ y 2 ∂ x 2 ] = [ I d 0 ∂ s ∂ x 1 ⊗ x 2 + ∂ t ∂ x 1 diag ( s ) ] \frac{\partial \mathbf{y}}{\partial \mathbf{x}} =\left[\begin{array}{ c c } \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}{ c c } \mathbb{I}_{d} & 0\\ \frac{\partial \mathbf{s}}{\partial \mathbf{x}_{1}} \otimes \mathbf{x}_{2} +\frac{\partial t}{\partial \mathbf{x}_{1}} & \operatorname{diag}( s) \end{array}\right] ∂x∂y=[∂x1∂y1∂x1∂y2∂x2∂y1∂x2∂y2]=[Id∂x1∂s⊗x2+∂x1∂t0diag(s)]
其中之所以是对角矩阵是因为
( x 1 ⋮ x n ) ⊙ ( y 1 ⋮ y n ) = ( x 1 ⋯ 0 ⋮ ⋱ ⋮ 0 ⋯ x n ) ( y 1 ⋮ y n ) \left(\begin{array}{ c } x_{1}\\ \vdots \\ x_{n} \end{array}\right) \odot \left(\begin{array}{ c } y_{1}\\ \vdots \\ y_{n} \end{array}\right) =\left(\begin{array}{ c c c } x_{1} & \cdots & 0\\ \vdots & \ddots & \vdots \\ 0 & \cdots & x_{n} \end{array}\right)\left(\begin{array}{ c } y_{1}\\ \vdots \\ y_{n} \end{array}\right) ⎝⎜⎛x1⋮xn⎠⎟⎞⊙⎝⎜⎛y1⋮yn⎠⎟⎞=⎝⎜⎛x1⋮0⋯⋱⋯0⋮xn⎠⎟⎞⎝⎜⎛y1⋮yn⎠⎟⎞
所以
∂ y 2 ∂ x 2 = ∂diag ( s ) x 2 ∂ x 2 = diag ( s ) \frac{\partial y_{2}}{\partial x_{2}} =\frac{\operatorname{\partial diag}(\mathbf{s}) x_{2}}{\partial x_{2}} =\operatorname{diag}(\mathbf{s}) ∂x2∂y2=∂x2∂diag(s)x2=diag(s)
这里一般会约束s大于0,所以一般神经网络输出的是log(s),然后取指数变回来。
Glow 一种可逆的1x1卷积
我们刚才随机划分的方法感觉处理图片的时候怪怪的,而且就算是随机交换channel也感觉不太对,有没有更优雅的方法?Glow给出了解决的方法,我们可以引入1x1可逆卷积核来代替这个划分的操作。其实1x1卷积核本身就有随机置换的味道在里面,只不过这篇文章的贡献在于用了一个trick保证了他的可逆性。
可以看个小例子(引用来自苏剑林的博客:https://kexue.fm/archives/5807),随机置换的操作其实就是一个简单的线性变换:
( 2 1 4 3 ) = ( 0 1 0 0 1 0 0 0 0 0 0 1 0 0 1 0 ) ( 1 2 3 4 ) \begin{pmatrix} 2\\ 1\\ 4\\ 3 \end{pmatrix} =\begin{pmatrix} 0 & 1 & 0 & 0\\ 1 & 0 & 0 & 0\\ 0 & 0 & 0 & 1\\ 0 & 0 & 1 & 0 \end{pmatrix}\begin{pmatrix} 1\\ 2\\ 3\\ 4 \end{pmatrix} ⎝⎜⎜⎛2143⎠⎟⎟⎞=⎝⎜⎜⎛0100100000010010⎠⎟⎟⎞⎝⎜⎜⎛1234⎠⎟⎟⎞
我们知道一个卷积核其实是可以表达成一个矩阵乘积的,只要我们把所有channel展开成一条向量,就可以写出一个矩阵的计算公式:
Y = X W Y=XW Y=XW
比如, x \displaystyle x x是 h ∗ w ∗ c \displaystyle h*w*c h∗w∗c的张量,c表示channel,对于1x1卷积来说,W就是一个 c ∗ c \displaystyle c*c c∗c的矩阵,就是他们的乘积,实际上可以将x想象成 h ∗ w \displaystyle h*w h∗w行 c \displaystyle c c列的矩阵,然后乘以W。
所以接下来的问题只有一个,如何保证这个W是可逆的。为了构造一个一定可逆的矩阵W,我们可以利用LU分解。因为任意矩阵都可以表达成
W = P L U W=PLU W=PLU
其中P是置换矩阵, L是下三角矩阵,对角线元素全为1,U是上三角矩阵,所以为了保证矩阵W可逆,那么只要保证P, L,U满秩就可以了,又因为P,L一定是满秩的,所以只要保证U满秩即可。那么一个方便的方法就是:
W = P L ( U + d i a g ( s ) ) W=PL( U+diag( s)) W=PL(U+diag(s))
其中,U是严格上三角矩阵,其对角线为0,我们只需保证这个s不为0即可。于是最终我们的可逆卷积核其导数行列式的求解为:
log ∣ det ( d conv 2 D ( h ; W ) d h ) ∣ = h ⋅ w ⋅ log ∣ det ( W ) ∣ = h ⋅ w ⋅ log ∣ d i a g ( s ) ∣ = h ⋅ w ⋅ s u m ( l o g ( ∣ s ∣ ) ) \log\left| \operatorname{det}\left(\frac{d\operatorname{conv} 2\mathrm{D} (\mathbf{h} ;\mathbf{W} )}{d\mathbf{h}}\right)\right| =h\cdot w\cdot \log |\operatorname{det} (\mathbf{W} )|\\ =h\cdot w\cdot \log |diag( s) |=h\cdot w\cdot sum( log( |s|)) log∣∣∣∣det(dhdconv2D(h;W))∣∣∣∣=h⋅w⋅log∣det(W)∣=h⋅w⋅log∣diag(s)∣=h⋅w⋅sum(log(∣s∣))
这里有个绝对值是因为这个分布变换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