在变分推断中，通常使用随机采样方法来估计模型的边际似然。这种过程中，每次采样产生的损失会随机波动；因此，验证损失也会随着波动。解决此问题的一种方法是通过增加采样次数来降低波动。下面是一个示例代码：

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

tfd = tfp.distributions

def build_vi_model(input_shape):
    model = tf.keras.Sequential([
        tf.keras.layers.Flatten(input_shape=input_shape),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(10),
        tfp.layers.DistributionLambda(lambda t: tfd.MultivariateNormalDiag(loc=t, scale_diag=[1]*10))
    ])
    return model

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.astype(np.float32) / 255.
x_test = x_test.astype(np.float32) / 255.

vi_model = build_vi_model(x_train.shape[1:])
opt = tf.keras.optimizers.Adam()

# 增加采样次数以降低波动
n_samples = 16

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        # 从模型中采样 n_samples 次
        y_pred = vi_model(x)
        y_pred_samples = y_pred.sample(n_samples)
        # 计算每个样本的边际似然并取平均
        log_likelihoods = y_pred_samples.log_prob(y[:, tf.newaxis, :])
        avg_log_likelihood = tf.reduce_mean(tf.reduce_sum(log_likelihoods, axis=-1), axis=0)
        # 计算 KL 散度
        kl_divergence = tfd.kl_divergence(y_pred.distribution, tfd.MultivariateNormalDiag(loc=[0]*10, scale_diag=[