单变量非参数密度模型实现.md

作者：Quan Chen

单变量非参数密度模型实现

最近在查阅图像压缩算法相关的论文《 Variational Image Compression with A Scale Hyperprior 》发现了关于单变量非参数密度模型的实现方式，今天就刚好使用tensorflow-keras API实现一下，有错误的地方,还望各大佬指正。

我们知道，随机变量的分布函数满足 ： $$ c(-\infty)=0 ; \quad c(\infty)=1 ; \quad p(x)=\frac{\partial c(x)}{\partial x} \geq 0 $$ 我们假设分布函数由一系列函数组成，然后通过链式法则求出密度函数 : $$ \begin{array}{l}{c=f_{K}\left(f_{K-1}\left(\ldots f_{1}()\right)\right)} \ {p=f_{K}^{\prime} \cdot f_{K-1}^{\prime} \ldots f_{1}^{\prime}}\end{array} $$ 通常来说， $f^{'}{k}$ 是雅克比矩阵， $\cdot$ 代表矩阵乘法。为了保证p(x)是单变量的。$f{1}$的域和$f_{K}$的范围应该都是一维的($d_k =r_K =1$)

为了保证p(x)是一个密度，我们只需要 $f_K$ 映射到0和1之间的范围，并确保 $p(x) \geq 0$。 为此，我们要求所有雅可比元素都是非负的。然后，计算p(x)的矩阵乘积也是非负的，我们已经定义了有效密度。 $$ \begin{array}{c}{f_{k}(x)=g_{k}\left(\boldsymbol{H}^{(k)} \boldsymbol{x}+\boldsymbol{b}^{(k)} \right),其中:1 \leq k<K} \ {f_{K}(\boldsymbol{x})=\operatorname{sigmoid}\left(\boldsymbol{H}^{(K)} \boldsymbol{x}+\boldsymbol{b}^{(K)}\right)} \ {g_{k}(\boldsymbol{x})=\boldsymbol{x}+\boldsymbol{a}^{(k)} \odot \tanh (\boldsymbol{x})}\end{array} $$ $a^k$是一个向量，⊙表示逐元素乘法，这种特殊非线性背后的基本原理是它允许扩展或收缩x=0附近的空间。当a大于0时它控制了膨胀率，当它小于0时控制收缩率。

密度函数表示为： $$ \begin{aligned} f_{k}^{\prime}(\boldsymbol{x}) &=\operatorname{diag} g_{k}^{\prime}\left(\boldsymbol{H}^{(k)} \boldsymbol{x}+\boldsymbol{b}^{(k)}\right) \cdot \boldsymbol{H}^{(k)} & 1 \leq k<K, \text { with } \ g_{k}^{\prime}(\boldsymbol{x}) &=1+\boldsymbol{a}^{(k)} \odot \tanh ^{\prime}(\boldsymbol{x}) & & \text { and } \ f_{K}^{\prime}(\boldsymbol{x}) &=\operatorname{sigmoid}^{\prime}\left(\boldsymbol{H}^{(K)} \boldsymbol{x}+\boldsymbol{b}^{(K)}\right) \cdot \boldsymbol{H}^{(K)} & & \end{aligned} $$ 由于密度函数是非负的，我们需要限制H为全非负的元素，同时a的元素需要以-1为下界，这可以通过再参数化实现： $$ \begin{aligned} \boldsymbol{H}^{(k)} &=\operatorname{softplus}\left(\hat{\boldsymbol{H}}^{(k)}\right) \ \boldsymbol{a}^{(k)} &=\tanh \left(\hat{a}^{(k)}\right) \end{aligned} $$

特殊情况下，设置K=1会产生一个逻辑分布：

$$ \begin{array}{l}{c(x)=\text { sigmoid }(h x+b)} \ {p(x)=\frac{h}{2} \cdot \frac{1}{1+\cosh (h x+b)}}\end{array} $$

代码实现：

from tensorflow.python.keras import layers
import tensorflow as tf

class NonParametric(layers.Layer):

    def __init__(self, K=1,
                 trainable=True,
                 name=None,
                 **kwargs):
        super().__init__(trainable=trainable,
                         name=name,
                         **kwargs)
        self._K = K
        self._K_value()  # 判断K是否大于0

    def _K_value(self):
        assert self._K > 0

    def build(self, input_shape):
        self._num_units = input_shape[-1].value

        self._H = []
        self._b = []

        if self._K == 1:
            self._H = self.add_variable("H", [self._num_units, self._num_units], self.dtype,
                                        initializer=tf.random_normal_initializer())
            self._b = self.add_variable("b", [self._num_units], self.dtype, initializer=tf.random_normal_initializer())

        else:
            self._a = []
            for i in range(self._K):
                self._H.append(
                    self.add_variable("H_" + str(i), [self._num_units, self._num_units], self.dtype,
                                      initializer=tf.random_normal_initializer())

                )
                self._b.append(
                    self.add_variable("b_" + str(i), [self._num_units], self.dtype,
                                      initializer=tf.random_normal_initializer())
                )

                self._a.append(
                    self.add_variable("a_" + str(i), [self._num_units], self.dtype,
                                      initializer=tf.random_normal_initializer())
                )

    def call(self, inputs: tf.Tensor, **kwargs):

        assert inputs.shape.ndims == 2
        self._H = tf.nn.softplus(self._H)

        def g_k(x: tf.Tensor, a_k) -> tf.Tensor:
            """g_k() 函数: x+ a^{k} + tanh(x)"""
            return x + a_k * tf.tanh(x)

        if self._K == 1:
            # 当K=1
            # f_K = sigmoid(H^{(K)}+b^{(b)}
            ys = tf.sigmoid(tf.matmul(inputs, self._H) + self._b)
            outputs = tf.gradients(ys, inputs)
            return outputs
        else:
            # 当K>1
            self._a = tf.tanh(self._a)
            # 1<= k <K
            f_k_input = inputs
            for i in range(self._K - 1):
                # f_k = g_k(H^{(k)}x+b^{(k)})
                # g_k(x) = x +a^{(k)}\odot \tanh ^{\prime}(\boldsymbol{x})
                f_k_input = g_k(tf.matmul(f_k_input, self._H[i]) + self._b[i], self._a[i])

            # k =K
            # f_K = sigmoid(H^{(K)}+b^{(b)}
            f_K = tf.sigmoid(tf.matmul(f_k_input, self._H[-1]) + self._b[-1])
            print("f_K", f_K)
            # 使用tensorflow自动求导
            outputs = tf.gradients(f_K, inputs)

            return outputs