Normalizing Flow流模型——NICE原理解析

1、前言

$\boxed{Normalizingflow}$，流模型，一种能够与目前流行的生成模型——$\boxed{\mathbf{GAN}、\mathbf{VAE}}$相媲美的模型。其也是一个生成模型，可是它的思路和另外两个的迂回策略却很大不同。本文我们就简单来介绍一个这个模型吧

论文：[1410.8516] NICE: Non-linear Independent Components Estimation (arxiv.org)

视频：[NICE——Normalizing Flow原理解析]

2、引入

在生成模型中，我们的目的就是计算出数据x的概率分布。

然而，数据的分布总是千奇百怪的。其$\boxed{无法被定义，无法被观测，无法被描述、无法被认知}$，说成是新时代的克苏鲁都不为过。

在GAN中，其简单粗暴，直接回避数据x的概率分布而采取其他策略；在VAE中，也是求的x的概率分布的下界。

可Normalizing flow就不一样了，这是一个真正“勇敢”的模型，它直面数据x的分布，并且确切的将其算出来了。

我们来看下面的模型图

假设x为数据。要找出它所服从的概率分布，但是由于太复杂，我们希望找到一个相对简单的概率的概率分布再加上一点其他简单的东西的表达它。如图，P(x)的分布很复杂，P(z1)相对来说没那么复杂，于是我们就想找到它们之间的关系式；而P(z1)虽然相对来说没那么复杂，但仍不是我们所能计算的。可是对于P(zn)，我们却可以计算出它，比如$P(zn)$就是一个标准的多元高斯分布。所以，理论上，我们只需要一步步往后递归，就可以得到P(x)的概率分布。

当随机变量时多维之时，实际上，概率分布之间雀食存在某种关系。我们

构造一个函数，$z=f(x)$，并且该函数存在反函数$x=f^{-1}(z)$。也就是x可以通过某个函数f转化成z（并且维度必须保持不变），则有
$$P_x(x)=P_z(z)\left|\det \frac{\partial z}{\partial x}\right|=P_z(f(x))\left|\det \frac{\partial f(x)}{\partial x}\right|=P_z(f(x))\left|\det J_f(x)\right|$$
其中det表示求里面矩阵的行列式，外面的||表示求绝对值。$J_f(x)$表示雅可比矩阵

该定理被称为Change of Variable Theorem

证明方法参考——随机变量的变量替换定理(Change of Variable Theorem)（此参考内容的表达形式与我所写存在一些差别，但是结论是一样的。我沿用了论文的写法）

3、目标函数

有了上面的转化式，我们就可以去定义目标函数了。一般地，对概率模型，我们就是采用对数极大似然估计的方法去估计出参数。所以（令$x_1=f_0(x_0)$，其中$x_0$表示原图像）
$$\begin{array}{rl}\log P_{x_0}(x_0)=& \log \left(P_{x_1}(x_1)\left|\det \frac{\partial x_1}{\partial x_0}\right|\right)\\ =& \log P_{x_1}(x_1)+\log \left|\det \frac{\partial x_1}{\partial x_0}\right|\\ =& \log \left(P_{x_2}(x_2)\left|\det \frac{\partial x_2}{\partial x_1}\right|\right)+\log \left|\det \frac{\partial x_1}{\partial x_0}\right|\\ =& \log P_{x_2}(x_2)+\log \left|\det \frac{\partial x_2}{\partial x_1}\right|+\log \left|\det \frac{\partial x_1}{\partial x_0}\right|\\ =& \log P_{x_2}(x_2)+\sum _{i=1}^2\log \left|\det \frac{\partial x_i}{\partial x_{i-1}}\right|\\ =& \cdots \cdots \\ =& \log P_{x_M}(x_M)+\sum _{i=1}^M\log \left|\det \frac{\partial x_i}{\partial x_{i-1}}\right|\end{array}\text{\hspace{1em}(1)}$$

其中，M表示有M层转化，所以，我们按照上面所提到的递归思想，取log之后变成连加，所以就变成了上面地式子，而$P_{x_M}(x_M)$表示的就是上面提到$\boxed{简单的，可以计算的概率分布}$。

$\boxed{目标很简单，最大化这个目标函数}$

所以，现在就有两个问题，

①第一个就是选择一个简单的先验分布$P_{x_M}$。

②第二个就是雅可比矩阵行列式要相对容易计算，不然假如我们的数据维度是1000维，那么它就是一个$1000\times 1000$维的矩阵了。计算行列式的计算量是相当大。

以上，就是Normalizing flow 的大致模型结构。下面我们就某一个具体的模型来讲解以下。

4、NICE

NICE（NON - LINEAR INDEPENDENT COMPONENTS ESTIMATION），非线性分量估计。

4.1、选择转化函数（或选择合适的雅可比矩阵）

4.1.1、分块耦合层

其思想是将x的维度按某种比例划分成两部分维度。假设x的维度是D维，那么就变成两部分$x_1\in (1:d)$，

$x_2\in (d+1:D)$。后面表示的是维度。

现在，对其进行函数变换得到$z_1,z_2$
$$\begin{gathered} z_1=x_1 \\ {z}_2=x_2+m(x_1) \end{gathered}$$
其中，$m(x)$是一个神经网络。最后将$z_1,z_2$堆叠起来，形成z。

在论文中，此处的$z_2$其实并不一定是这样，他表达为$z_2=g(x_2,m(x_1))$，$g$可以是其他函数。但论文中作者在实践的时候就是使用的加性耦合层。并且其具有特殊性。

我们知道，每一次的转化，我们都是需要计算其雅可比矩阵的行列式，作为目标函数的一部分。我们来看它的雅可比矩阵
$$\left[\begin{array}{cc}\frac{\partial z_1}{\partial x_1}& \frac{\partial z_1}{\partial x_2}\\ \frac{\partial z_2}{\partial x_1}& \frac{\partial z_2}{\partial x_2}\end{array}\right]=\left[\begin{array}{cc}1& 0\\ \frac{\partial z_2}{\partial x_1}& 1\end{array}\right]$$
我们计算并不需要矩阵本身，我们需要的是它的行列式，所以此处行列式的值不就是1吗？而前面所提到的$\log \left|\det \frac{\partial z}{\partial x}\right|=\log 1=0$

后面生成数据的时候是需要反函数的，所以有反函数。
$$\begin{gathered} x_1=z_1 \\ {x}_2=z_2-m(z_1) \end{gathered}$$
容易看到，我们对其进行分块，一部分经过了变化，一部分不经过变化。我们进行M次的变化，如果每次都是$x_1$不经过变化，只变化$x_2$，这样是不合理的。因此，我们会交替进行，假设上面的z为第一次变化。在第二次变化，我们就
$$\begin{gathered} z_2^{(2)}=z_2 \\ {z}_1^{(2)}=z_1+m(z_2) \end{gathered}$$

4.1.2、缩放

在经过了M次的分块耦合层之后，论文中提到，会最终的输出$z^{(M)}$做一次缩放。即引入一个与$z_n$相同维度的向量s。记最终结果为h（最后的先验分布也暂时记为$P_h$），则
$$h=s\cdot z^{(M)}$$
即对应元素相乘。同样的，求出它的雅可比矩阵（仍然以二维为例）
$$\left[\begin{array}{cc}\frac{\partial h_1}{\partial z_1^{(M)}}& \frac{\partial h_1}{\partial z_2^{(M)}}\\ \frac{\partial h_2}{\partial z_1^{(M)}}& \frac{\partial h_2}{\partial z_2^{(M)}}\end{array}\right]=\left[\begin{array}{cc}{s}_1& 0\\ 0& s_2\end{array}\right]$$
所以行列式为（D为$s_i$的维度）
$$\sum _{i=1}^D\log s_i$$

其反函数为
$$z^{(M)}=s^{-1}h$$
所以对于目标函数可以变为
$$\log P_{x_0}(x_0)=\log P_h(h)+\sum _{i=1}^D\log s_i$$

4.2、选择合适的先验分布$P_h(h)$

在论文中，其对先验分布是假设各个维度都相互独立的。即
$$P(h)=\prod _{i=1}^{D}P(h_i)$$
并且对其分布论文是给了两个建议，一个是高斯分布，另一个就是logistic分布。

假如是高斯分布的时候（并且是标准高斯）
$$\begin{array}{rl}\log P(h)=& \log \prod _{i=1}^{D}P(h_i)\\ =& \sum _{i=1}^D\log P(h_i)\\ =& \sum _{i=1}^D\log \frac{1}{\sqrt{2\pi }}\exp \left\{-\frac{h_i^2}{2}\right\}\\ =& \sum _{i=1}^D\left(\log \frac{1}{\sqrt{2\pi }}-\frac{h_i^2}{2}\right)\end{array}$$
前面的$\log$那一项显然不在我们优化的参数之内，$\boxed{\mathbf{我们最终目标函数可以写成}}$
$$\log P_{x_0}(x_0)=\sum _{i=1}^D\left(\log s_i-\frac{h_i^2}{2}\right)$$

$\boxed{容易看到，当每一个维度对应的s_i越大，则说明该维度越不重要。因为如果是重要的s_i，那么h_i刚好与其相反，阻止其增大。}$

5、代码实现（Pytorch）

根据原论文，其有4个加性耦合层，一个加性耦合层里面的m(x)是五层1000个神经元的神经网络。对于Mnist这个数据集来说，选用的是Logistic分布。采样方法是从0，1分布中采样，通过反函数采样出z。论文中训练的很长，长达1500个epochs之后。我仅仅训练了1000个（用了整整两个小时，并且用GPU）。(感兴趣的可以训练到2000或更多，应当会好一些)

Ps：代码内容有所参考，但是时间久远，已经找不到原作者了。如有人知道，望告知，十分感谢！

import torch
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import transforms
from  torch import nn
from tqdm import tqdm
import matplotlib.pyplot as plt
from torch.functional import  F
class Coupling(nn.Module):
    def __init__(self,input_dim,hidden_dim,hidden_layer,odd_flag):
        '''
        加性耦合层
        @param input_dim: 输入维度
        @param hidden_dim: 隐藏层维度
        @param hidden_layer: 隐藏层个数
        @param odd_flag: 当前耦合层是否是在整个模型中属于奇数（用作调换切分顺序）
        '''
        super().__init__()
        #用作判断是否需要互换切分的位置
        self.odd_flag=odd_flag%2

        #五层隐藏层，神经元1000
        self.input_transform=nn.Sequential(
            nn.Linear(input_dim//2,hidden_dim),
            nn.ReLU()
        )
        self.m=nn.ModuleList([
            nn.Sequential(
                nn.Linear(hidden_dim,hidden_dim),
                nn.ReLU(),
            ) for _ in range(hidden_layer-1)
        ])
        #输出为原始维度
        self.out_transform=nn.Sequential(
            nn.Linear(hidden_dim,input_dim//2),
        )
    def forward(self,x,reverse):
        '''
        @param x:  数据 [ batch_size , 784 ]
        @param reverse: 是否是反向推导（z->x）
        @return:
        '''
        batch_size,W=x.shape #取出维度
        #重构维度（为了切分）
        x=x.reshape(batch_size,W//2,2)

        #按奇偶性划分
        if self.odd_flag:
            x1, x2 = x[:, :, 0], x[:, :, 1]
        else:
            x2, x1 = x[:, :, 0], x[:, :, 1]

        #将x2输入神经网络
        input_transfrom=self.input_transform(x2)
        for i in self.m:
            input_transfrom=i(input_transfrom)
        out_transform=self.out_transform(input_transfrom)

        #是否是反向推导
        if reverse:
            x1=x1-out_transform #反函数
        else:
            x1=x1+out_transform

        #将数据组合回来
        if self.odd_flag:
            x=torch.stack((x1,x2),dim=2)
        else:
            x=torch.stack((x2,x1),dim=2)

        return x.reshape(-1,784)

class Scale(nn.Module):
    def __init__(self,input_dim):
        '''
        缩放层
        @param input_dim: 输入数据维度
        '''
        super().__init__()

        #构造与数据同维度的s
        self.s=nn.Parameter(torch.zeros(1,input_dim))
    def forward(self,x,reverse):
        '''
        @param x: 输入数据
        @param reverse: 是否是反向推导
        @return:
        '''
        if reverse:
            result=torch.exp(-self.s)*x #反函数
        else:
            result=torch.exp(self.s)*x
        return result,self.s
class NICE(nn.Module):
    def __init__(self,couping_num):
        '''
        @param couping_num: 耦合层个数
        '''
        super().__init__()
        #初始化耦合层
        self.couping=nn.ModuleList([
            Coupling(784,1000,5,odd_flag=i+1)
            for i in torch.arange(couping_num)
        ])
        #初始化缩放层
        self.scale=Scale(784)

    def forward(self,x,reverse):

        '''
        前向推导
        @param x: 输入数据
        @param reverse: #是否是反向
        @return:
        '''
        for i in self.couping:
            x=i(x,reverse)
        h,s=self.scale(x,reverse)
        return h,s
    def likeihood(self,h,s):
        #计算极大似然估计
        loss_s = torch.sum(s) #s的log雅可比行列式损失
        log_prob = proior.log_prob(h) #logictic分布极大似然
        loss_prob = torch.sum(log_prob, dim=1) #按列求和
        loss = loss_s + loss_prob #总损失
        #由于pytorch是最小值优化，故取反
        return -loss
    def generate(self,h):
        '''
        @param h: logistic分布采样所得
        @return:
        '''
        z,s=self.scale(h,True)
        for i in reversed(self.couping):
            z=i(z,True)
        return z
def train():
    #归一化
    transformer = transforms.Compose([
        transforms.ToTensor()
    ])
    #载入数据
    data = MNIST("data", transform=transformer, download=True)
    #存入写入器
    dataloader = DataLoader(data, batch_size=200, shuffle=True,num_workers=4)
    #初始化模型
    nice = NICE(4).to(device)
    #优化器
    optimer = torch.optim.Adam(params=nice.parameters(), lr=1e-3,eps=1e-4,betas=(0.9,0.999))

    #开始训练
    epochs = 1000

    for epoch in torch.arange(epochs):
        loss_all = 0
        dataloader_len = len(dataloader)
        for i in tqdm(dataloader, desc="第{}轮次".format(epoch)):
            sample, label = i
            sample = sample.reshape(-1, 784).to(device)

            h, s = nice(sample, False) #预测
            loss = nice.likeihood(h, s).mean() #计算损失

            optimer.zero_grad() #归零
            loss.backward() #反向传播
            optimer.step() #更新

            with torch.no_grad():
                loss_all += loss
        print("损失为{}".format(loss_all / dataloader_len))
        torch.save(nice, "nice.pth")

class Logistic(torch.distributions.Distribution):
    '''
    Logistic 分布
    '''
    def __init__(self):
        super().__init__()

    def log_prob(self, x):

        return -(F.softplus(x) + F.softplus(-x))

    def sample(self, size):

        z = torch.distributions.Uniform(0., 1.).sample(size)
        return torch.log(z) - torch.log(1. - z)
if __name__ == '__main__':
    # 是否有闲置GPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    #先验分布
    proior = Logistic()

    #训练
    train()

    #预测
    x=proior.sample((10,784))#采样

    #载入模型
    nice=torch.load("nice.pth",map_location=device)
    #生成数据
    result=nice.generate(x)
    result=result.reshape(-1,28,28)
    for i in range(10):
        plt.subplot(2,5,i+1)
        img=result[i].detach().numpy()
        plt.imshow(img)
        plt.gray()
    plt.show()

6、结束

以上，就是Normalizing flow 的全部内容了。如有问题，还望指出。

7、参考

①随机变量的变量替换定理(Change of Variable Theorem) - 郑之杰的个人网站 (0809zheng.github.io)