混合密集网络（Mixture Density Networks）

前言

考古了1994年的一篇文章，写的很朴实，不像现在很多的AI文章有一种过度包装的感觉，论文题目《Mixture Density Networks》。

理论

混合密集网络是一种将高斯混合模型和神经网络结合的网络，与一般的神经网络不同，它在网络的输出部分不再使用线性层或softmax作为预测值，为了引入模型的不确定性，认为每个输出是一种高斯混合分布，而不是一个确定值或者单纯的高斯分布，至于为什么不是通过高斯分布引入不确定性，这也是混合密集网络的一大亮点：高斯混合分布可以解决高斯分布不好解决的多值映射问题。
设输入为$\mathbf{x}$，输出为$\mathbf{t}$，这里以回归问题为例，输出和输出均是可能有多个维度的矢量。目标值的概率密度可以表示成多个核函数的线性组合：
$$p(\mathbf{t}|\mathbf{x})=\sum _{i=1}^m{\alpha }_i(\mathbf{x}){\varphi }_i(\mathbf{t}|\mathbf{x})$$
其中${\alpha }_i(\mathbf{x})$称为混合系数，可以认为是$\mathbf{x}$的一种先验概率，${\varphi }_i$表示目标向量$\mathbf{t}$的第$i$个核。这里的核函数可以有多种选择，文章选了最经常使用的形式：高斯分布，因为足够数量的混合高斯分布在原理上可以近似任何一个分布。$m$是高斯混合分布选取了多少个核，核函数$\varphi$表示成：
$${\varphi }_i(\mathbf{t}|\mathbf{x})=\frac{1}{(2\pi {)}^{c/2}{\sigma }_i(\mathbf{x}{)}^c}\exp \{-\frac{||\mathbf{t}-{\mathbf{\mu }}_i(\mathbf{x})|{|}^2}{2{\sigma }_i(\mathbf{x}{)}^2}\}$$
其中$c$为$\mathbf{t}$的维度。注意这里的每一个核函数都是一个多元的高斯分布，${\sigma }_i(\mathbf{x})$是一个标量，但是${\mathbf{\mu }}_i(\mathbf{x})$是一个与目标值$\mathbf{t}$同维度的矢量，反应回归的预测值。除此之外，与认为$\mathbf{t}$的不同元之间是相互独立的高斯分布不同，高斯混合分布不需要这个假设。

如何搭建网络

网络的搭建也非常简单，前面是神经网络，在输出的时候建立混合高斯模型，重点在怎么连接上，我们先来看高斯混合模型需要多少个参数值：
根据$p(\mathbf{t}|\mathbf{x})$的公式，${\alpha }_i(\mathbf{x}),{\sigma }_i(\mathbf{x}),{\mathbf{\mu }}_i(\mathbf{x})$都属于要优化的参数，所以从神经网络传出来的参数量应该等于高斯混合模型要优化的参数量，一个$p(\mathbf{t}|\mathbf{x})$有$m$个${\alpha }_i(\mathbf{x})$，$m$个${\sigma }_i(\mathbf{x})$，${\mathbf{\mu }}_i(\mathbf{x})$中有$cm$个标量，所以对应的神经网络的输出应该是$(c+2)m$个输出变量。
在高斯混合分布中，所有的混合系数之和为1：
$$\sum _{i=1}^m{\alpha }_i(\mathbf{x})=1$$
在神经网络里面可以通过softmax函数实现：
$${\alpha }_i=\frac{\exp (z_i^{\alpha })}{{\sum }_{j=1}^M\exp (z_j^{\alpha })}$$
其中$z^{\alpha }$对应神经网络的一个输出变量，相应地，每个高斯单元的方差和均值也可以表示为：
$$\begin{gathered} {\sigma }_i=\exp (z_i^{\sigma }) \\ {\mu }_{ik}=z_{ik}^{\mu } \end{gathered}$$
网络的损失函数是寻找以给定的$\mathbf{x}$为条件下什么样的参数可以使$p(\mathbf{t}|\mathbf{x})$的概率最大，即 $\mathrm{argmax}p(\mathbf{t}|\mathbf{x})$，通常我们写成误差函数的形式：
$$\begin{gathered} \mathcal{L}=\sum _q{\mathcal{L}}^q \\ {\mathcal{L}}^q=-\ln \{\sum _{i=1}^m{\alpha }_i({\mathbf{x}}^q){\varphi }_i({\mathbf{t}}^q|{\mathbf{x}}^q\} \end{gathered}$$
其中${\mathcal{L}}^q$表示每个样本的损失。根据高斯混合分布中哪一个成分占的多，该成分对应的中心值${\mathbf{\mu }}_i$即为该样本的预测值，更详细一点，在所有的成分中：
$$\underset{i}{\max }\{\frac{{\alpha }_i(\mathbf{x})}{{\sigma }_i(\mathbf{x}{)}^c}\}$$
其中这个$i$对应的${\mathbf{\mu }}_i$就是预测值。但是这个预测值是一个近似的值，在严格意义上，预测值应该是计算得到的分布$p(\mathbf{t}|\mathbf{x})$关于$\mathbf{t}$的条件平均（也就是分布函数关于$\mathbf{t}$的期望）：
$$\begin{gathered} E(t)=\sum _i{\alpha }_i(\mathbf{x})\int \mathbf{t}{\varphi }_i(\mathbf{t}|\mathbf{x})d\mathbf{t} \\ =\sum _i{\alpha }_i(\mathbf{x}){\mathbf{\mu }}_i(\mathbf{x}) \end{gathered}$$
模型的不确定性仍然通过方差来计算：
$$\begin{gathered} s^2(\mathbf{x})=E[\mathbf{t}-E(\mathbf{t}){]}^2 \\ =\sum _i{\alpha }_i(\mathbf{x})\{{\sigma }_i(\mathbf{x}{)}^2+||{\mathbf{\mu }}_i(\mathbf{x})-\sum _j{\alpha }_j(\mathbf{x}){\mathbf{\mu }}_j(\mathbf{x})|{|}^2\} \end{gathered}$$

假设预测值服从高斯混合分布有什么好处

传统的神经网络模型对于单值映射具有良好的拟合能力，但是对于一个输入可能存在多个输出的情况，拟合效果就会很差，比如一个从$t$到$x$的映射函数；
$$x=t+0.3\sin (2\pi t)+ϵ$$
对于一个输入$t$，只有一个$x$与之对应，此时拟合效果如下：

忽略图中的坐标含义，横坐标表示输入，纵坐标表示输出。
但是，如果现在对调$t$和$x$的位置，当有一个输入的时候，可能存在多个输出，这样一来，网络的拟合能力就会大幅降低：

其中曲线为神经网络的拟合曲线，散点为样本实际分布。
这个问题出现的原因是我们采用真实值-预测值的均平方作为损失函数，网络自己在学习的时候只保证这一个目标，但是很明显全局的均方误差最小并不能保证每个样本都有合适的拟合值，其实这个问题在图像生成选取损失函数的时候也会出现，当我们选取整幅图像所有像素点的均方误差作为损失函数优化网络的时候，虽然损失在减小，但是生成的图像可能是模糊的，这是全局最优不能代表局部最优的一个典型案例。
而假设输出为高斯混合分布就可以很好地解决这个问题，当我们假设输出服从单一的高斯分布时，其实是在默认输出只有一个可能值，这个值就是高斯分布的峰值对应的横坐标，但是高斯混合分布含有多个成分${\alpha }_i$，不同成分的动态大小就是就输出可能值的一种反应，作者用一幅图很好的说明了不同成分对预测值的影响：

可以看到，在样本$\mathbf{x}$的预测值只有一个的情况下（如$\mathbf{x}=0.2,\mathbf{x}=0.8$），$p(\mathbf{t}|\mathbf{x})$中只有一个成分占有很高的比重，并且相应的该成分对应的${\alpha }_i$的值接近于1，而对应有多个可能值的情况下（如$\mathbf{x}=0.5$），$p(\mathbf{t}|\mathbf{x})$中的三个成分占有的比重不相上下，每个成分对应的峰值也都差不多。这也引出了一个问题，如何选择总的成分的个数呢？也就是说：$p(\mathbf{t}|\mathbf{x})$对应的$m$应该是多少？
文章给出的建议是$m$的个数应该大于等于最大的样本可能的预测值的个数，上面的图也是对这种选择的一种解释：潜在预测值的个数可以通过$p(\mathbf{t}|\mathbf{x})$里面不同${\alpha }_i$所占的比重体现。

程序实现

模型部分

"""A module for a mixture density network layer

For more info on MDNs, see _Mixture Desity Networks_ by Bishop, 1994.
"""
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from torch.distributions import Categorical
import math


ONEOVERSQRT2PI = 1.0 / math.sqrt(2 * math.pi)

# 标注一下输出维度out_features = 1，对应论文从神经网络出来的维度是（c + 2） * m = (1 + 2) * 5,但是mdn里面用了三个linear层来分别表示pi，sigma和mu
class MDN(nn.Module):
    """A mixture density network layer

    The input maps to the parameters of a MoG probability distribution, where
    each Gaussian has O dimensions and diagonal covariance.

    Arguments:
        in_features (int): the number of dimensions in the input
        out_features (int): the number of dimensions in the output
        num_gaussians (int): the number of Gaussians per output dimensions

    Input:
        minibatch (BxD): B is the batch size and D is the number of input
            dimensions.

    Output:
        (pi, sigma, mu) (BxG, BxGxO, BxGxO): B is the batch size, G is the
            number of Gaussians, and O is the number of dimensions for each
            Gaussian. Pi is a multinomial distribution of the Gaussians. Sigma
            is the standard deviation of each Gaussian. Mu is the mean of each
            Gaussian.
    """

    def __init__(self, in_features, out_features, num_gaussians):
        super(MDN, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.num_gaussians = num_gaussians
        self.pi = nn.Sequential(
            nn.Linear(in_features, num_gaussians),
            nn.Softmax(dim=1)
        )
        self.sigma = nn.Linear(in_features, out_features * num_gaussians)
        self.mu = nn.Linear(in_features, out_features * num_gaussians)

    def forward(self, minibatch):
        pi = self.pi(minibatch)  # [btz, num_gaussians]
        sigma = torch.exp(self.sigma(minibatch))  # [btz, num_gaussians]
        # 因为sigma和mu都与输出的维度有关，所以在这里还要展开，给输出一个维度
        sigma = sigma.view(-1, self.num_gaussians, self.out_features)
        mu = self.mu(minibatch)
        mu = mu.view(-1, self.num_gaussians, self.out_features)
        return pi, sigma, mu


def gaussian_probability(sigma, mu, target):
    """Returns the probability of `target` given MoG parameters `sigma` and `mu`.
    example: sigma: torch.Size([150, 5, 1]) mu: torch.Size([150, 5, 1]) target: [150, 1]
    Arguments:
        sigma (BxGxO): The standard deviation of the Gaussians. B is the batch
            size, G is the number of Gaussians, and O is the number of
            dimensions per Gaussian.
        mu (BxGxO): The means of the Gaussians. B is the batch size, G is the
            number of Gaussians, and O is the number of dimensions per Gaussian.
        target (BxI): A batch of target. B is the batch size and I is the number of
            input dimensions.

    Returns:
        probabilities (BxG): The probability of each point in the probability
            of the distribution in the corresponding sigma/mu index.
        返回高斯混合分布的component：phi,如果输出是多维的有 exp(a) * exp(b) = exp(a+b)
    """
    target = target.unsqueeze(1).expand_as(sigma)
    ret = ONEOVERSQRT2PI * torch.exp(-0.5 * ((target - mu) / sigma)**2) / sigma
    return torch.prod(ret, 2)


def mdn_loss(pi, sigma, mu, target):
    """Calculates the error, given the MoG parameters and the target
        pi: torch.Size([150, 5]) sigma: torch.Size([150, 5, 1]) mu: torch.Size([150, 5, 1])
    The loss is the negative log likelihood of the data given the MoG
    parameters.
    """
    prob = pi * gaussian_probability(sigma, mu, target)
    nll = -torch.log(torch.sum(prob, dim=1))
    return torch.mean(nll)


def sample(pi, sigma, mu):
    """Draw samples from a MoG.
    """
    # Choose which gaussian we'll sample from,返回采样点的索引
    # 返回的是 均值 + 方差*随机噪声 的形式
    pis = Categorical(pi).sample().view(pi.size(0), 1, 1)
    # Choose a random sample, one randn for batch X output dims
    # Do a (output dims)X(batch size) tensor here, so the broadcast works in
    # the next step, but we have to transpose back.
    gaussian_noise = torch.randn(                                        # [2, 150]
        (sigma.size(2), sigma.size(0)), requires_grad=False)
    # torch.gather(dim=1) 表示按照列号进行索引，寻找采样的pi对应的sigma
    variance_samples = sigma.gather(1, pis).detach().squeeze()   # [150]
    mean_samples = mu.detach().gather(1, pis).squeeze()
    return (gaussian_noise * variance_samples + mean_samples).transpose(0, 1)

测试部分

"""A script that shows how to use the MDN. It's a simple MDN with a single
nonlinearity that's trained to output 1D samples given a 2D input.
"""
import matplotlib.pyplot as plt
import sys
sys.path.append('../mdn')
from MDN.mdn import mdn
import torch
import torch.nn as nn
import torch.optim as optim

# 输入为2维的向量，输出为一个标量，高斯分布的成分有5个
input_dims = 2
output_dims = 1
num_gaussians = 5


def translate_cluster(cluster, dim, amount):
    """Translates a cluster in a particular dimension by some amount
    torch.add_:
                一般来说函数加了下划线的属于内建函数，将要改变原来的值，没有加下划线的并不会改变原来的数据，引用时需要另外赋值给其他变量
    """
    translation = torch.ones(cluster.size(0)) * amount
    cluster.transpose(0, 1)[dim].add_(translation)
    return cluster


print("Generating training data... ", end='')
cluster1 = torch.randn((50, input_dims + output_dims)) / 4
cluster1 = translate_cluster(cluster1, 1, 1.2)
cluster2 = torch.randn((50, input_dims + output_dims)) / 4
cluster2 = translate_cluster(cluster2, 0, -1.2)
cluster3 = torch.randn((50, input_dims + output_dims)) / 4
cluster3 = translate_cluster(cluster3, 2, -1.2)
training_set = torch.cat([cluster1, cluster2, cluster3]) # torch.Size([150, 3])
print('Done')


print("Initializing model... ", end='')
model = nn.Sequential(
    nn.Linear(input_dims, 5),
    nn.Tanh(),
    mdn.MDN(5, output_dims, num_gaussians)
)
optimizer = optim.Adam(model.parameters())
print('Done')

print('Training model... ', end='')
sys.stdout.flush()
# training_set的前两列作为训练数据，后一列作为预测值，对应in_features和out_features
for epoch in range(1000):
    model.zero_grad()
    pi, sigma, mu = model(training_set[:, 0:input_dims])
    loss = mdn.mdn_loss(pi, sigma, mu, training_set[:, input_dims:])
    loss.backward()
    optimizer.step()
    if epoch % 100 == 99:
        print(f' {round(epoch/10)}%', end='')
        sys.stdout.flush()
print('Done')

# 这一步骤用来计算预测值，实际上均值已经可以表示预测值，作者在这里加上了方差*随机噪声用来表示模型的不确定性
print('Generating samples... ', end='')
pi, sigma, mu = model(training_set[:, 0:input_dims])
samples = mdn.sample(pi, sigma, mu)
print('Done')

print('Saving samples.png... ', end='')
fig = plt.figure()
ax = fig.add_subplot(projection='3d')

xs = training_set[:, 0]
ys = training_set[:, 1]
zs = training_set[:, 2]

ax.scatter(xs, ys, zs, label='target')
ax.scatter(xs, ys, samples, label='samples')
ax.legend()
fig.savefig('samples.png')
print('Done')