Почему в этом случае градиент вычисляется с помощью GradientTape неправильно (с использованием tfp.vi.monte_carlo_variational_loss)

Question

Градиенты от tf.GradientTape, похоже, не соответствуют правильному минимуму в функции, которую я пытаюсь минимизировать.

Я пытаюсь использовать вариационный вывод чёрного ящика tenorflowprobability (используя tf2), сtf.GradientTape, оптимизатор keras, вызывающий функцию apply_gradients. Суррогатный апостериор является простым 1d Normal. Я пытаюсь приблизить пару нормалей, см. Функцию pdist. Для простоты я просто пытаюсь оптимизировать параметр масштаба.

Текущий код:

from scipy.special import erf
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd

def pdist(x):
    return (.5/np.sqrt(2*np.pi)) * np.exp((-(x+3)**2)/2) + (.5/np.sqrt(2*np.pi)) * np.exp((-(x-3)**2)/2)
def logpdist(x):
    logp = np.log(1e-30+pdist(x))
    assert np.all(np.isfinite(logp))
    return logp

optimizer = tf.keras.optimizers.Adam(learning_rate=0.1)
mu = tf.Variable(0.0,dtype=tf.float64)
scale = tf.Variable(1.0,dtype=tf.float64)
for it in range(100):
    with tf.GradientTape() as tape:
        surrogate_posterior = tfd.Normal(mu,scale)
        elbo_loss = tfp.vi.monte_carlo_variational_loss(logpdist,surrogate_posterior,sample_size=10000)
    gradients = tape.gradient(elbo_loss, [scale])
    optimizer.apply_gradients(zip(gradients, [scale]))
    if it%10==0: print(scale.numpy(),gradients[0].numpy(),elbo_loss.numpy())

Вывод (показывая каждую 10-ю итерацию):

SCALE   GRAD   ELBO_LOSS
1.100, -1.000, 2.697
2.059, -0.508, 1.183
2.903, -0.354, 0.859 <<< (right answer about here)
3.636, -0.280, 1.208
4.283, -0.237, 1.989
4.869, -0.208, 3.021
5.411, -0.187, 4.310
5.923, -0.170, 5.525
6.413, -0.157, 7.250
6.885, -0.146, 8.775

По некоторым причинамГрадиент не отражает истинный градиент, который должен быть около нуля в масштабе = 2,74.

Почему градиент не относится к фактическому elbo_loss?

Answer 1

Будем надеяться, что кто-то может объяснить, почему предыдущая реализация не удалась (а также почему она ничего не делает, но вместо этого просто имеет неправильный ответ). Во всяком случае, я обнаружил, что могу исправить это, убедившись, что ключевые выражения используют библиотеку математических тензорных потоков, а не библиотеку numpy. В частности, замена двух методов выше на

def pdist(x):
    return (.5/np.sqrt(2*np.pi)) * tf.exp((-(x+3)**2)/2) + (.5/np.sqrt(2*np.pi)) * tf.exp((-(x-3)**2)/2)

def logpdist(x):
    return tf.math.log(pdist(x))

Стохастическая оптимизация теперь работает.

Вывод:

2.020, -0.874, 1.177
2.399, -0.393, 0.916
2.662, -0.089, 0.857
2.761, 0.019, 0.850
2.765, 0.022, 0.843
2.745, -0.006, 0.851
2.741, 0.017, 0.845
2.752, 0.005, 0.852
2.744, 0.015, 0.852
2.747, 0.013, 0.862

Я не собираюсь принимать мой собственный ответя был бы признателен, если бы можно было дать несколько ответов, дающих представление о том, почему это сейчас работает и почему раньше это не удавалось (и почему режим отказа не был исключением или подобным, а вместо этого неправильным градиентом).