在张量流概率中，如何更新仅在KL散度中使用的可学习先验？

Question

我正在研究一种可变自动编码器，我希望对潜伏分布的KL-散度正则化中使用的先验方法进行loc（均值）和小数位数（stddev）的更新。

下面的代码段是一个人为设计的最小示例，它演示了我正在尝试实现的目标。这开始起作用，但是在一些随机的时期之后（有时为1，有时为200，但通常为7或8），然后只是freezes。没有错误消息或其他任何内容。

loc = tf.Variable(tf.random.normal([ndim], stddev=0.1, dtype=tf.float32))
scale = tfp.util.TransformedVariable(
    tf.random.normal([ndim], mean=1.0, stddev=0.1, dtype=tf.float32),
    bijector=tfb.Chain([tfb.Shift(1e-5), tfb.Softplus(), tfb.Shift(0.5413)]))
prior = tfd.Independent(tfd.Normal(loc=loc, scale=scale), reinterpreted_batch_ndims=1)

_input = tfkl.Input(shape=(1,))
_loc = tfkl.Dense(ndim, name="loc_params")(_input)
_scale = tfkl.Dense(ndim, name="untransformed_scale_params")(_input)
_scale = tf.math.softplus(_scale + np.log(np.exp(1) - 1)) + 1e-5
_output = tfpl.DistributionLambda(
    make_distribution_fn=lambda t: tfd.Independent(tfd.Normal(loc=t[0], scale=t[1])),
    activity_regularizer=tfpl.KLDivergenceRegularizer(prior, use_exact_kl=True, weight=0.1)
)([_loc, _scale])
model = tf.keras.Model(_input, _output)
model.compile(optimizer='adam', loss=lambda y_true, model_out: -model_out.log_prob(y_true))
hist = model.fit(ds, epochs=N_EPOCHS, verbose=2)

我有可运行的要点here。

一个更具体的示例，以及与我尝试更新和简化的体系结构接近的结构，是disentangled_vae的tfp示例。在其手动训练循环中，使用持久性tf.Variables参数化了新的tfd.MultivariateNormalDiag is instantiated on every loop。我正在尽力避免手动训练循环，并且还尝试使用更多类似Keras的语法，因此我不想直接移植此示例。

非常感谢任何建议。谢谢！

---

编辑：activity_regularizer在附加到潜在（瓶颈）分布时似乎工作正常。我有一个更完整的示例in this Colab notebook。由于这在我的体系结构中有效，因此我不再需要答案。

但是，我强烈怀疑冻结模型拟合是理想的行为，因此这仍然是一个问题。

Answer 1

由于机械在大多数情况下都可以工作，而不仅仅是上面的人为冻结示例，我不再认为这是需要答案的问题。

我已经通过tensorflow-probability存储库问题页面报告了无错误的冻结行为。 See here。