Как рассчитать KL-дивергенцию для распределения продукта с помощью Tensorflow

Question

Я хочу вычислить дивергенцию KL D KL (∏q (zk | x) ‖ p (z)) в тензорном потоке

---

, где ∏q (zk | x) является произведением N распределений (от k = 1 до k = N), не зависящих друг от друга, типа tenorflow_probability.layers.MultivariateNormalTriL. p (z) является tfp.distributions.MultivariateNormalDiag. Расхождение KL между одиночным q (z | x) и p (z): D KL (q (z | x) ‖ p (z)) работает.

---

Я попытался подкласс класса Распределения tfp, чтобы создать распределение продукта Q (z) = ∏q (zk | x) и затем вычислить D KL (Q (z) ‖ p (z)), но я должен был реализовать функцию _kl_divergence, это вернуло меня к началу.

---

Я хотел бы решить это с вероятностью тензорного потока или, по крайней мере, тензорным потоком.

Answer 1

Я не думаю, что есть аналитика c KL для этого случая. Но поскольку KL [q || p] равно q (x) * log_q (x) -q (x) * log_p (x), вы можете использовать приближение Монте-Карло для KL. В TFP это выглядит примерно так:

q = QDist(...)
p = MVNDiag(...)
x = q.sample()  # x ~ q(x)
kl = q.log_prob(x) - p.log_prob(x)

или, для более низкой дисперсии

x = q.sample(10)  # x ~ q(x)
kl = tf.math.reduce_mean(q.log_prob(x) - p.log_prob(x))