使用神经网络进行分布变换

过去的文章采样：从均匀分布采样、正太分布采样到MCMC提及过分布变换，即从一个分布变换到另外一个分布，例如从均匀分布变换到正太分布。不过还真没有尝试过直接使用神经网络进行分布变换机器应用。今天这里进行一个简单的尝试和实践。

分布变换

假设有神经网络$f$，分布变换要做的事情是把从满足某一分布的数据$X$，变换为数据$Y$，

$Y = f(X)$

使其满足目标分布。

一个具体的例子，假如有数据$[X_{1}, X_{2}, \dots, X_{n}]$独立地采样字均匀分布，神经网络$f$负责变换为$[Y_{1}, Y_{2}, \dots, Y_{n}]$使其满足正太分布。

均值估计，

$\mu = \frac{1}{n} \sum_{i=1}^{n} Y_{i}$

方差的无偏估计，

$\sigma^{2} = \frac{1}{n-1} \sum_{i=1}^{n} (Y_{i}-\mu)^{2}$

通过KL散度来评估与标准正太分布的差距，计算KL散度，

$\begin{aligned} \operatorname{KL}\Big(N(\mu,\sigma^2)\Big\Vert N(0,1)\Big) =&\int \frac{1}{\sqrt{2\pi\sigma^2}}e^{-(x-\mu)^2/2\sigma^2} \left(\log \frac{e^{-(x-\mu)^2/2\sigma^2}/\sqrt{2\pi\sigma^2}}{e^{-x^2/2}/\sqrt{2\pi}}\right)dx\\ =&\int \frac{1}{\sqrt{2\pi\sigma^2}}e^{-(x-\mu)^2/2\sigma^2} \log \left\{\frac{1}{\sqrt{\sigma^2}}\exp\left\{\frac{1}{2}\big[x^2-(x-\mu)^2/\sigma^2\big]\right\} \right\}dx\\ =&\frac{1}{2}\int \frac{1}{\sqrt{2\pi\sigma^2}}e^{-(x-\mu)^2/2\sigma^2} \Big[-\log \sigma^2+x^2-(x-\mu)^2/\sigma^2 \Big] dx \\ =& \frac{1}{2}\Big[-\log \sigma^{2}+\mu^{2}+\sigma^{2}-1\Big] \end{aligned}$

因此有，

$\operatorname{KL}\Big(N(\mu_{i},\sigma_{i}^{2})\Big\Vert N(0,1)\Big) =\frac{1}{2}\Big(-\log \sigma_{i}^{2}+\mu_{i}^{2}+\sigma_{i}^{2}-1\Big)$

上式作为loss参与神经网络的训练，主导参数的优化方向。

更一般的形式，

$\operatorname{KL} \Big( N(\mu_{1},\sigma_{1}^{2})\,\|\,N(\mu_{2},\sigma_{2}^{2}) \Big)= \frac {1}{2}\left(-\log {\frac {\sigma _{1}^{2}}{\sigma _{2}^{2}}} + \frac {(\mu _{1}-\mu _{2})^{2}}{\sigma _{2}^{2}} + {\frac {\sigma _{1}^{2}}{\sigma _{2}^{2}}}-1\right)$

可以用来控制生成的正太分布的形状。

拉普拉斯分布的概率密度，

$f(x\mid \mu ,b)={\frac {1}{2b}}\exp \left(-{\frac {|x-\mu |}{b}}\right)\,$

均值和方差分别为$\mu, 2b^{2}$。

实验

上述方法最关键是的KL散度这一层的实现，

class KLLossLayer(tf.keras.layers.Layer):

    def __init__(self, **kwargs):
        super(KLLossLayer, self).__init__(**kwargs)

    def call(self, inputs):
        hdims = tf.cast(tf.shape(inputs)[-1], tf.float32)
        mu = tf.reduce_mean(inputs, axis=1, keepdims=True)
        var = tf.reduce_sum(tf.square(inputs - mu), axis=1) / (hdims-1)
        kl_loss = 0.5 * tf.reduce_mean(-tf.math.log(var) + tf.square(mu) + var - 1)
        self.add_loss(kl_loss)
        return inputs

模型十分简单，

inputs = tf.keras.layers.Input(shape=(hdims,))
x = inputs
x = tf.keras.layers.Dense(hdims, activation="tanh")(x)
x = tf.keras.layers.Dense(hdims)(x)
outputs = KLLossLayer()(x)
model = tf.keras.Model(inputs, outputs)
model.compile(optimizer="adam")

变换结果，