如何计算两个张量流概率分布的 Kullback-Leibler 散度相对于分布均值的梯度？

Question

在 tensorflow-2.0 中，我正在尝试创建一个 keras.layers.Layer，它输出两个 tensorflow_probability.distributions 之间的 Kullback-Leibler (KL) 散度。我想计算输出的梯度（即 KL 散度）相对于 tensorflow_probability.distributions.

之一的平均值

到目前为止，在我所有的尝试中，结果梯度是 0，不幸的是。

我尝试实现如下所示的最小示例。我想知道这些问题是否与 tf 2 的急切执行模式有关，因为我知道在 tf 1 中使用了类似的方法，默认情况下禁用急切执行。

这是我尝试过的最小示例：

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Layer,Input

# 1 Define Layer

class test_layer(Layer):

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

    def build(self, input_shape):
        self.mean_W = self.add_weight('mean_W',trainable=True)

        self.kernel_dist = tfp.distributions.MultivariateNormalDiag(
            loc=self.mean_W,
            scale_diag=(1.,)
        )
        super(test_layer, self).build(input_shape)

    def call(self,x):
        return tfp.distributions.kl_divergence(
            self.kernel_dist,
            tfp.distributions.MultivariateNormalDiag(
                loc=self.mean_W*0.,
                scale_diag=(1.,)
            )
        )

# 2 Create model

x = Input(shape=(3,))
fx = test_layer()(x)
test_model = Model(name='test_random', inputs=[x], outputs=[fx])


# 3 Calculate gradient

print('\n\n\nCalculating gradients: ')

# example data, only used as a dummy
x_data = np.random.rand(99,3).astype(np.float32)

for x_now in np.split(x_data,3):
#     print(x_now.shape)
    with tf.GradientTape() as tape:
        fx_now = test_model(x_now)
        grads = tape.gradient(
            fx_now,
            test_model.trainable_variables,
        )
        print('\nKL-Divergence: ', fx_now, '\nGradient: ',grads,'\n')

print(test_model.summary())

上面代码的输出是

Calculating gradients: 

KL-Divergence:  tf.Tensor(0.0029436834, shape=(), dtype=float32) 
Gradient:  [<tf.Tensor: id=237, shape=(), dtype=float32, numpy=0.0>] 


KL-Divergence:  tf.Tensor(0.0029436834, shape=(), dtype=float32) 
Gradient:  [<tf.Tensor: id=358, shape=(), dtype=float32, numpy=0.0>] 


KL-Divergence:  tf.Tensor(0.0029436834, shape=(), dtype=float32) 
Gradient:  [<tf.Tensor: id=479, shape=(), dtype=float32, numpy=0.0>] 

Model: "test_random"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 3)]               0         
_________________________________________________________________
test_layer_3 (test_layer)    ()                        1         
=================================================================
Total params: 1
Trainable params: 1
Non-trainable params: 0
_________________________________________________________________
None

KL散度计算正确，但得到的梯度为0。获得梯度的正确方法是什么？

Answer 1

如果有人感兴趣，我发现了如何解决这个问题：

行

self.kernel_dist = tfp.distributions.MultivariateNormalDiag(
            loc=self.mean_W,
            scale_diag=(1.,)
        )

不应该在层 class 定义的 build()- 方法内部，而是在 call() 方法内部。这是修改后的例子：

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Layer,Input

# 1 Define Layer

class test_layer(Layer):

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

    def build(self, input_shape):
        self.mean_W = self.add_weight('mean_W',trainable=True)
        super(test_layer, self).build(input_shape)

    def call(self,x):
        self.kernel_dist = tfp.distributions.MultivariateNormalDiag(
            loc=self.mean_W,
            scale_diag=(1.,)
        )
        return tfp.distributions.kl_divergence(
            self.kernel_dist,
            tfp.distributions.MultivariateNormalDiag(
                loc=self.mean_W*0.,
                scale_diag=(1.,)
            )
        )

# 2 Create model

x = Input(shape=(3,))
fx = test_layer()(x)
test_model = Model(name='test_random', inputs=[x], outputs=[fx])


# 3 Calculate gradient

print('\n\n\nCalculating gradients: ')

# example data, only used as a dummy
x_data = np.random.rand(99,3).astype(np.float32)

for x_now in np.split(x_data,3):
#     print(x_now.shape)
    with tf.GradientTape() as tape:
        fx_now = test_model(x_now)
        grads = tape.gradient(
            fx_now,
            test_model.trainable_variables,
        )
        print('\nKL-Divergence: ', fx_now, '\nGradient: ',grads,'\n')

print(test_model.summary())

现在的输出是



Calculating gradients: 

KL-Divergence:  tf.Tensor(0.024875917, shape=(), dtype=float32) 
Gradient:  [<tf.Tensor: id=742, shape=(), dtype=float32, numpy=0.22305119>] 


KL-Divergence:  tf.Tensor(0.024875917, shape=(), dtype=float32) 
Gradient:  [<tf.Tensor: id=901, shape=(), dtype=float32, numpy=0.22305119>] 


KL-Divergence:  tf.Tensor(0.024875917, shape=(), dtype=float32) 
Gradient:  [<tf.Tensor: id=1060, shape=(), dtype=float32, numpy=0.22305119>] 

Model: "test_random"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_2 (InputLayer)         [(None, 3)]               0         
_________________________________________________________________
test_layer_1 (test_layer)    ()                        1         
=================================================================
Total params: 1
Trainable params: 1
Non-trainable params: 0
_________________________________________________________________
None

符合预期。

这是从 tensorflow 1 更改为 tensorflow 2 的内容吗？

Answer 2

我们正在通过分布和双射器工作，使它们能够友好地关闭构造函数中的变量。（尚未完成 MVN。）与此同时，您可以使用 tfd.Independent(tfd.Normal(loc=self.mean_W, scale=1), reinterpreted_batch_ndims=1)，我认为它可以在您的构建方法中使用，因为我们已经采用了 Normal.

另外：你见过tfp.layers包吗？特别是 https://www.tensorflow.org/probability/api_docs/python/tfp/layers/KLDivergenceAddLoss 您可能会感兴趣。

如何计算两个张量流概率分布的 Kullback-Leibler 散度相对于分布均值的梯度？

How to calculate the gradient of the Kullback-Leibler divergence of two tensorflow-probability distributions with respect to the distribution's mean?

tensorflow

keras

tensorflow-probability