DUL学习心得（二）：flow模型之RealNVP

        在学习伯克利CS294-158-SP20第3节课时，课程中提到的一种flow模型的结构RealNVP,并在课后作业也有相关的练习，于是，笔者读了这篇论文，并对课程中的基本知识进行理解，然后跟着课后作业，分析了代码。论文在此，如有兴趣请查看。

        Flow 模型是基于样本进行概率密度估计变通方法中的一种，其他的变通方法还有VAE，GAN。为什么会有出现这些变通方法呢？因为基于样本算它们所服从的最优分布的太难算了，于是大家就想办法简化它或者近似计算它。

Real NVP是flow模型的一种，它的目标仍然是降低计算复杂度而尽量不损失精度。

flow模型的基本概念

改变变量

     ——让神经网络去学习这个复杂的计算吧

原本计算公式应该是：

其中z是怎么来的呢？

我们要对x的概率密度做估计，就是计算概率的积分

然而连续空间的积分计算很麻烦，按照神经网络的精神，难以计算或者难以总结出规律的，我们就设置几个隐藏层，让隐藏层来帮我们寻找计算规律。于是，概率密度估计就成了这样

训练神经网络的目标是

将输出用一个隐变量替换，则变成

为什么这里要替换成一个隐变量，并且，将隐变量设定为服从某个分布呢？因为我们的初衷是预测x的概率密度做估计，也就是希望从现有的一堆数据样本X中去估计这堆数据它服从怎样的分布，为什么要去预测它服从怎样的分布呢？因为这堆数据样本是离散的，当我们要想生成符合这堆数据样本规律的其他更多的，甚至近乎于连续的数据时，我们需要知道这堆数据符合的规律，也就是它服从的最好的分布。当我们知道一个一个分布了，要生成数据，只需要按照要求从这个分布中采样即可。

而假定x服从某个分布太过生硬，也很难计算，于是，就引入隐变量z，并且

这样，就把预测x的分布转换为估计z的分布。

其中，

因为，z本质上还是由x生成的， 那么

这就需要f(x)是可导且单调。

因此训练目标可以替换为

Real NVP的主要思想

上面的训练目标需要计算f的导数的范数，导数的话，就是计算雅克比矩阵，如果x的维度很大的话，计算量会非常大。因此，减少计算最直接的办法是让化简导数的行列式。什么样的行列式计算起来最容易呢？当然是对角矩阵。对角矩阵的行列式就是对角线上各个元素相乘。因此，RealNVP的主要思想是通过精心设计一个神经网络，让输出对应输入的导数行列式可以转化为类似于对角线的矩阵。

基于这个目标，作者设计了一个双目标网络，每个layer是一个仿射变换，仿射变换让输出对应输入的导数行列式转换为对角线相乘，达到简化计算的效果。多个layer堆叠，让每一维输入能被均衡的处理。

Affine Transform

仿射变换的原理是讲输入分为两块，第一块不做任何计算直接输出，另一块输出采用自回归的方式，基于第一块计算输出。有一点残差计算的感觉。

具体来说，一个仿射变换的计算过程如下：

 

其中s() function 基于第一块计算第二块输入的尺度缩放系数，t() function 对第一块输入做变换，在后面的程序中，s和t由mlp完成。 是哈达玛乘积，即相应位置的元素相乘。

 

经过这样的变换，由于第一块的输出等于输入，导数为1， 也就是一个单位矩阵，同时，第一块的输出和第二块输入不相关，导数是零矩阵，第二块对滴一块的导数比较复杂，先不管。根据公式，右下角，第二块输出对应第二块输入的导数就是s function的自然对数。整个导数公式如下：

这是个上三角矩阵，矩阵的行列式是单位矩阵和为对角线的乘积，由于s function和第二块输入做哈达玛乘积，s function的输出在和第二块做乘积的时候应该转为与输入同维度的对角线矩阵。整个导数行列式就变得很容易计算了。

将公式转换为流程图，是这样的

从输出的分布中采样生成数据时，正好是反过来的

Mask

仿射变换中，输入数据分块通过掩码来完成，直接输出的部分被掩去，掩码用向量b来表示，仿射变换表示为

                        

多层结构

上面的仿射变换虽然计算简单，但是x1:d在单个flow中是不会改变的，如果每一个flow都按照这种方式，那整个flow模型，每一层输出的前半部分是完全一样的。整个模型携带和处理信息的能力就会变差。为了克服这种问题，整个flow采用交替mask，也就是上一层是（1：d)被掩码的话，这一层就换做(d:D)掩码，下一层再换成(1:d)，如此交替下去，让输出更加均衡。

代码分析

class MLP(nn.Module):
    def __init__(self, input_size, n_hidden, hidden_size, output_size):
        super().__init__()
        layers = []
        for _ in range(n_hidden):
            layers.append(nn.Linear(input_size, hidden_size))
            layers.append(nn.ReLU())
            input_size = hidden_size
        layers.append(nn.Linear(hidden_size, output_size))
        self.layers = nn.Sequential(*layers)
            
    def forward(self, x):
        return self.layers(x)


class AffineTransform(nn.Module):
    def __init__(self, type, input_size=2, n_hidden=2, hidden_size=256):
        super().__init__()
        """
        input_size： 数据的维度
        n_hidden: mlp的层数
        hidden_size: 每层隐层单元个数
        b：对应于公式中的b。
        scale和scale_shfit暂时没搞懂出处
        """

        self.mask = self.build_mask(type=type).to(ptu.device)
        self.scale = nn.Parameter(torch.zeros(1), requires_grad=True).to(ptu.device)
        self.scale_shift = nn.Parameter(torch.zeros(1), requires_grad=True).to(ptu.device)
        # input 
        self.mlp = MLP(input_size=input_size, n_hidden=n_hidden, hidden_size=hidden_size, output_size=2).to(ptu.device)

    def build_mask(self, type):
        """
        制作mask模板，将每层的x1（这里不是指前一维，而是指直接输出，并作为x2自回归前一项的那一半）掩去。
        直接输出的部分值为1，需要计算的部分值为0，原因在于，最后输出是 y = b*x+(1-b)y'
        """
        # if type == "left", left half is a one
        # if type == right", right half is a one
        assert type in {"left", "right"}
        if type == "left":
            mask = ptu.FloatTensor([1.0, 0.0])
        elif type == "right":
            mask = ptu.FloatTensor([0.0, 1.0])
        else:
            raise NotImplementedError
        return mask

    def forward(self, x, reverse=False):
        """
        核心计算部分
        输出包含两个部分，y(该层网络结构的输出）和 导数的行列式（也是s_function的输出）
        """
        
        def forward(self, x, reverse=False):
        # returns transform(x), log_det
        batch_size = x.shape[0]
        
        # 对输入进行掩码
        # x_ = b * x
        mask = self.mask.repeat(batch_size, 1)
        x_ = x * mask

        ## 通过mlp来模拟S function和t function
        # 由于最后导数的行列式做对数计算之后才是对数似然的第二项，这里mlp直接模拟log(det(df/dx))
        log_s, t = self.mlp(x_).split(1, dim=1) #todo: fix
        
        # log_s 在s的基础上再做了一次线性变换，暂时没看懂为什么
        log_s = self.scale * torch.tanh(log_s) + self.scale_shift
        
        # t = (1-b) * t_function(b*x)
        t = t * (1.0 - mask)
        
        # log_s = （1-b）* s_function((b*x))
        log_s = log_s * (1.0 - mask)

        if reverse:  # inverting the transformation
            x = (x - t) * torch.exp(-log_s)
        else:
            # 程序对应的公式： y = x * exp((1-b)*(s'* s_function(b*x) + s_bais) + (1-b) * t_function(b*x)
            # 原公式： y = (b*x) + (1-b)*x*exp(s_function(b*x)) + (1-b)* t_function(b*x)
            x = x * torch.exp(log_s) + t

        return x, log_s


class RealNVP(nn.Module):
    def __init__(self, 
                 transforms):
        super().__init__()
        
        # 用高斯分布作为新样本生成的隐变量采样分布，相当于style transfer中以白噪声作为输入生成新的照片。
        self.prior = torch.distributions.Normal(torch.tensor(0.).to(ptu.device), torch.tensor(1.).to(ptu.device))
        self.transforms = nn.ModuleList(transforms)
        
    def flow(self, x):
        """
        输出隐变量z和导数的行列式
        """
        # maps x -> z, and returns the log determinant (not reduced)
        z, log_det = x, torch.zeros_like(x)
        for op in self.transforms:
            z, delta_log_det = op.forward(z)
            log_det += delta_log_det
        return z, log_det

    def invert_flow(self, z):
        """用于抽样的反向计算"""
        # z -> x (inverse of f)
        for op in reversed(self.transforms):
            z, _ = op.forward(z, reverse=True)
        return z

    def log_prob(self, x):
        """
        对数似然：sum( log Pz(f(x)) + log_det ) 
        """
        z, log_det = self.flow(x)
        return torch.sum(log_det, dim=1) + torch.sum(self.prior.log_prob(z), dim=1)

    def sample(self, num_samples):
        """
        用于生成过程中的采样
        """
        z = self.prior.sample([num_samples, 2])
        return self.invert_flow(z)

    def nll(self, x):
        """最终的loss"""
        return - self.log_prob(x).mean()

# 调用程序
# 多个仿射层堆叠，并且掩码是交替进行的，也就是直接输出部分是交替着来的
# n_hidden, hidden_size根据需求定
real_nvp = RealNVP([AffineTransform("left", n_hidden=2, hidden_size=64),
                    AffineTransform("right", n_hidden=2, hidden_size=64),
                    AffineTransform("left", n_hidden=2, hidden_size=64),
                    AffineTransform("right", n_hidden=2, hidden_size=64),
                    AffineTransform("left", n_hidden=2, hidden_size=64),
                    AffineTransform("right", n_hidden=2, hidden_size=64)])

# loss计算
loss = real_nvp.nll(x)

# 隐变量计算
latents = real_nvp.log_prob(x)

# 新样本的生成
new_samples = real_nvp.sample(8)