i-ResNet:通过变分解量化和结构设计改进流模型.

1. 流模型

流模型(flow-based model)是一种从观测变量$x$到简单隐变量$z$的可逆变换$z=f(x)$,该变换可以通过叠加多个简单的可逆变换构造$f(x) = f_1 ◦ \cdots ◦ f_L(x)$,由于每个变换$f_i$可逆且容易求Jacobian行列式,因此采样过程也容易实现$f^{-1}(z) = f^{-1}_L ◦ \cdots ◦ f^{-1}_1(z)$。

根据概率密度的变量替换公式和链式法则,观测变量$x$的分布为:

\[p(x) = p(z)\cdot |\det \prod_{i=1}^{L} \frac{\partial f_i}{\partial f_{i-1}}|\]

其对数似然函数为:

\[\log p(x) = \log p(z) + \sum_{i=1}^{L} \log |\det \frac{\partial f_i}{\partial f_{i-1}}|\]

此时不需要显式地计算分布$p(x)$的概率密度函数,而是通过初始分布$p(z)$的概率密度(常取标准正态分布)以及映射过程产生的Jacobian行列式计算即可。

由于必须保证逆变换简单和Jacobian行列式容易计算,则每次变换$f_i$的非线性变换能力都很弱。为了保证充分的拟合能力,模型必须堆得非常深,此时参数量和计算量非常大。

2. Invertible ResNet

本文作者在ResNet结构基础上增加了一些约束,使得模型可逆,同时保留了ResNet的基本结构和拟合能力。

下图为标准ResNet与可逆ResNet的对比图。可逆ResNet允许信息无损可逆流动,而标准ResNet在某处则存在信息瓶颈的“坍缩”现象。

ResNet的设计思路是用神经网络$g(\cdot)$拟合输出与输入的残差:

\[y = x + g(x)\]

(1)什么时候可逆?

上式可逆的一个充分条件是:

\[\text{Lip} (g) = \mathop{\max}_{x_1 \ne x_2} \frac{||g(x_1)-g(x_2)||_2}{||x_1-x_2||_2} \lt 1\]

即函数$g(\cdot)$的Lipschitz范数小于1。通常函数$g(\cdot)$是由神经网络实现的,而神经网络是由矩阵运算和激活函数组合而成的:

\[g(x) = \sigma(Wx+b)\]

若希望函数$g(\cdot)$的Lipschitz范数小于1,则应有激活函数$\sigma$的Lipschitz范数不超过1且$Wx+b$的Lipschitz范数小于1。

激活函数$\sigma$是标量函数,其Lipschitz范数为导数值,常用的激活函数均满足导数不超过1。

$Wx+b$的Lipschitz范数小于1即矩阵$W$的谱范数小于1。可以对其做谱归一化然后乘以0~1之间的一个常数$W \leftarrow cW/ ||W||_2$。

(2)逆函数是什么?

若$y=x+g(x)$可逆,其逆函数为$x=h(y)$。则问题转变为求解非线性方程组。考虑如下迭代:

\[x_{n+1} = y -g(x_n)\]

则迭代序列\(\{x_n\}\)是$y$的函数,如果\(\{x_n\}\)收敛到固定函数:

\[\mathop{\lim}_{n \to \infty} x_{n}(y) = \hat{h}(y)\]

此时有$\hat{h}(y) = y - g(\hat{h}(y))$,$\hat{h}(y)$即为所求逆函数。因此如果上述迭代过程收敛,则收敛结果即为所求逆函数。

根据可逆的Lipschitz条件:

\[\forall x_1,x_2 \quad ||g(x_1)-g(x_2)||_2 \leq \text{Lip} (g) ||x_1-x_2||_2\]

有如下关系:

\[\begin{aligned} ||x_{n+1}-x_n||_2 &= ||g(x_n)-g(x_{n-1})||_2 \leq \text{Lip} (g) ||x_n-x_{n-1}||_2 \\ &= \text{Lip} (g) ||g(x_{n-1})-g(x_{n-2})||_2 \leq \text{Lip} (g)^2 ||x_{n-1}-x_{n-2}||_2 \\ & \cdots \\ & \leq \text{Lip} (g)^n ||x_{1}-x_{0}||_2 \end{aligned}\]

因此序列\(\{x_n\}\)收敛的一个充分条件是$\text{Lip} (g) < 1$,这与可逆条件是吻合的。且逆函数为$x_{n+1} = y -g(x_n)$的不动点。数值计算时,只需要迭代一定步数,使得满足精度要求即可。

(3)如何计算Jacobian行列式

变换$y=x+g(x)$的Jacobian矩阵计算为:

\[J_y = \frac{\partial }{\partial x} (x+g(x)) = I+ \frac{\partial g}{\partial x} = I+J_g\]

优化目标需要计算Jacobian行列式的对数值:

\[\log |\det(J_y)| = \log |\det(I+J_g)| = \log \det(I+J_g)\]

根据恒等式$\log \det (A) = \text{Trace}(\log (A))$,上式可以展开为级数:

\[\log \det(I+J_g) = \text{Trace}(\log (I+J_g))= \text{Trace}(\sum_{n=1}^{\infty} (-1)^{n-1}\frac{J_g^n}{n})\]

该级数收敛的条件是$||J_g||_2 < 1$,即$\text{Lip} (g) < 1$,这与可逆条件也是吻合的。对该无穷级数截断$N$项,则有:

\[\text{Trace}(\log (I+J_g))= \sum_{n=1}^{N} (-1)^{n-1}\frac{\text{Trace}(J_g^n)}{n} + \mathcal{O}(\text{Lip} (g)^N)\]

上式把矩阵$\log (I+J_g)$的迹的计算转换为矩阵$J_g^n$的迹的计算,需要计算矩阵的$n$次方。考虑$p(u)$是一个多元概率分布,其均值为$0$、协方差为单位矩阵,则对于任意矩阵$A$,有

\[\text{Trace}(A) = \Bbb{E}_{u \text{~} p(u)}[u^TAu]\]

作者提出,对于每次迭代,从$p(u)$中只随机选择一个向量$u$,用$u^TAu$作为$\text{Trace}(A)$的近似值:

\[\text{Trace}(\log (I+J_g)) ≈ \sum_{n=1}^{N} (-1)^{n-1}\frac{u^TJ_g^nu}{n} , \quad u \text{~} p(u)\]

矩阵$u^TJ_g^nu$的计算可以拆解成多次矩阵与向量的乘法$u^TJ_g(\cdots (J_g(J_g u)))$,从而避开直接计算$J_g^n$。

3. 实验分析

作者首先设计了一个toy数据集,人为构造一些有规律的随机点,然后用生成模型去拟合它的分布。结果表明可逆ResNet产生了对称的、没有偏置的结果,而Glow的结果则是有偏。这可能是因为Glow需要以某种方式打乱输入并对半切分,对半之后两部分的运算是不一样的,这就带来了不对称。

作者测试了使用可逆ResNet的分类结果。在不同的$c=\text{Lip} (g)$数值下分类结果如下表:

作者也给出了一些图像生成结果: