【深度学习】Normalizing flow原理推导+Pytorch实现

1、前言

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

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转化成z（并且维度必须保持不变）
P x ( x ) = P z ( z ) ∣ det ⁡ ∂ z ∂ x ∣ = P z ( f ( x ) ) ∣ det ⁡ ∂ f ( x ) ∂ x ∣ 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| Px​(x)=Pz​(z) ​det∂x∂z​ ​=Pz​(f(x)) ​det∂x∂f(x)​ ​
其中det表示求里面矩阵的行列式，外面的||表示求绝对值。

$\boxed{证明：}$
$${\int }_x{P}_x(x)dx={\int }_z{P}_z(z)dz=1$$
对不定积分，始终有
$$|P_x(x)dx|=|P_z(z)dz|$$
因为概率P始终大于0，所以，有
P x ( x ) ∣ d x ∣ = P z ( z ) ∣ d z ∣ ↓ P x ( x ) = P z ( z ) ∣ d z d x ∣ ↓ P x ( x ) = P z ( f ( x ) ) ∣ d f ( x ) d x ∣ \begin{aligned} P_x(x)|dx|=&P_z(z)|dz| \\\downarrow \\P_x(x)=&P_z(z)\left|\frac{dz}{dx}\right|\\\downarrow \\P_x(x)=&P_z(f(x))\left|\frac{df(x)}{dx}\right| \end{aligned} Px​(x)∣dx∣=↓Px​(x)=↓Px​(x)=​Pz​(z)∣dz∣Pz​(z) ​dxdz​ ​Pz​(f(x)) ​dxdf(x)​ ​​
当x，z的维度是高维时，$\frac{df(x)}{dx}$就变成了一个矩阵，但是概率怎么可能会去使用矩阵来表达呢？所以实际上里面是要加上一个行列式，得
P x ( x ) = P z ( f ( x ) ) ∣ det ⁡ d f ( x ) d x ∣ P_x(x)=P_z(f(x))\left|\det\frac{df(x)}{dx}\right| Px​(x)=Pz​(f(x)) ​detdxdf(x)​ ​
对于$\frac{df(x)}{dx}$，其实它还有另一个名字——$\boxed{\mathbf{雅可比矩阵}}$。

假如$x=\left(\begin{array}{cc}{x}_1& x_2\end{array}\right)$是一个二维向量，那么雅可比矩阵可表示为(各个分量相互求导)
$$\frac{df(x)}{dx}=\left[\begin{array}{cc}\frac{d{f}_1(x)}{d{x}_1}& \frac{d{f}_2(x)}{d{x}_1}\\ \frac{d{f}_2(x)}{d{x}_1}& \frac{d{f}_2(x)}{d{x}_2}\end{array}\right]$$
值得注意的是，$\boxed{\mathbf{z}=\mathbf{f}(\mathbf{x})\mathbf{中的}\mathbf{f}必须存在反函数，即存在\mathbf{x}={\mathbf{f}}^{-1}(\mathbf{z})}$，因为后面我们生成样本的时候就是从P(z)中去生成的。要反着去计算。即

3、目标函数

有了上面的转化式，我们就可以去定义目标函数了。一般地，对概率模型，我们就是采用对数极大似然估计的方法去估计出参数。所以
log ⁡ P x ( x ) = log ⁡ ( P z ( f ( x ) ) ∣ det ⁡ d f ( x ) d x ∣ ) = log ⁡ P z ( f M ( x ) ) + ∑ k = 1 M log ⁡ ∣ det ⁡ d f k ( x ) d x k ∣ (1) \begin{aligned} \log P_x(x)=&\log \left(P_z(f(x))\left|\det\frac{df(x)}{dx}\right|\right) \\=&\log P_z(f_M(x))+ \sum\limits_{k=1}^M\log\left|\det \frac{df_k(x)}{dx_k}\right| \end{aligned}\tag{1} logPx​(x)==​log(Pz​(f(x)) ​detdxdf(x)​ ​)logPz​(fM​(x))+k=1∑M​log ​detdxk​dfk​(x)​ ​​(1)

其中，M表示有M层转化，对于行列式，我们有$|AB|=|A||B|$，所以，我们按照上面所提到的递归思想，实际上就是M个雅可比矩阵相乘，取log之后变成连加，所以就变成了上面地式子，而$P_z(f(x_m))$是实际上就是上面提到$\boxed{简单的，可以计算的概率分布}$。

当然了，上面表示的是某一个样本，那么对于所有的样本我们记作$X$

所以就变成了（单个样本用$x^i$表示，共有N个样本，并且样本与样本之间独立同分布）
$$\log P_X(X)=\log \prod _{i=1}^N{P}_x(x^i)=\sum _{i=1}^N\log P_x(x^i)$$
每一个$\log P_x{x}^i$就是（式1）里面所提到的，所以可以写成
log ⁡ P X ( X ) = ∑ i = 1 N ( log ⁡ P z ( f ( x M i ) ) + ∑ k = 1 M log ⁡ ∣ det ⁡ d f k ( x i ) d x k i ∣ ) \log P_X(X)=\sum\limits_{i=1}^N \left( \log P_z(f(x^i_M))+ \sum\limits_{k=1}^M\log\left|\det \frac{df_k(x^i)}{dx^i_k}\right| \right) logPX​(X)=i=1∑N​(logPz​(f(xMi​))+k=1∑M​log ​detdxki​dfk​(xi)​ ​)
$\boxed{目标很简单，最大化这个目标函数}$

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

①第一个就是选择一个简单的先验分布$P_z$。

②第二个就是雅可比矩阵行列式要相对容易计算，不然假如我们的数据维度是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{d{z}_2}{\partial x(x_1)}& 1\end{array}\right]$$
我们计算并不需要行列式本身，我们需要的是它的行列式，所以此处行列式的值不就是1吗？而前面所提到的 log ⁡ ∣ det ⁡ d f k ( x ) d x k ∣ = log ⁡ 1 = 0 \log\left|\det \frac{df_k(x)}{dx_k}\right|=\log 1 =0 log ​detdxk​dfk​(x)​ ​=log1=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，则
$$h=s\cdot z^{(M)}$$
即对应元素相乘。同样的，求出它的雅可比矩阵（仍然以二维为例）
[ ∂ h 1 ∂ z 1 ( M ) ∂ h 1 ∂ z 2 ( M ) ∂ h 2 ∂ z 1 ( M ) ∂ h 2 ∂ z 2 ( M ) ] = [ s 1 0 0 s 2 ] \begin{bmatrix} \frac{\partial h_1}{\partial z^{(M)}_1} & \frac{\partial h_1}{\partial z^{(M)}_2} \\ \frac{\partial h_2}{\partial z^{(M)}_1} & \frac{\partial h_2}{\partial z^{(M)}_2} \end{bmatrix}= \begin{bmatrix} s_1 & 0 \\ 0 & s_2 \end{bmatrix} ​∂z1(M)​∂h1​​∂z1(M)​∂h2​​​∂z2(M)​∂h1​​∂z2(M)​∂h2​​​ ​=[s1​0​0s2​​]
所以行列式为
$$\sum _{i=1}^D\log s_i$$
D为$s_i$的维度。

其反函数为
$$z^{(M)}=s^{(-1)}h$$
所以对于目标函数可以变为
$$\log P_X(X)=\sum _{i=1}^N\left(\log P_z(h)+\sum _{i=1}^D\log s_i\right)$$

4.2、选择合适的先验分布$P(z)$

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

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

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

5、代码实现（Pytorch）

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

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(self, self).__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 的全部内容了。如有问题，还望指出。