TensorFlow Probability中混合模型可训练权重参数配置问题

正确实现方案 #

你遇到的问题核心是两个：

1. tfd.Categorical的概率参数需要满足非负、总和为1的约束，无约束的独立变量不仅会导致概率分布非法，还会出现梯度传递断裂
2. 嵌套定义tf.Variable（将tf.Variable作为另一个tf.Variable的初始值）会导致内层变量脱离计算图，无法被训练追踪

推荐使用更稳定的logits参数定义混合权重，不需要手动做概率约束，梯度传递也完全正常，完整修改后的代码如下：

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd
import numpy as np

# 这里替换成你的真实数据
X = np.random.randn(272, 1)

# 定义混合权重的logits变量，形状对应混合成分数量，不需要手动约束和为1
weight_logits = tf.Variable([0.2, 0.8], name='weight_logits')

dist = tfd.Mixture(
    cat=tfd.Categorical(logits=weight_logits),  # 用logits参数代替probs
    components=[
        tfd.Normal(loc=tf.Variable(0., name='loc1'), scale=tf.Variable(1., name='scale1')),
        tfd.Normal(loc=tf.Variable(0., name='loc2'), scale=tf.Variable(1., name='scale2'))
    ]
)

optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)

@tf.function
def train_step(X):
    with tf.GradientTape() as tape:
        loss = -tf.reduce_mean(dist.log_prob(X)) 
    gradients = tape.gradient(loss, dist.trainable_variables)     
    optimizer.apply_gradients(zip(gradients, dist.trainable_variables)) 
    return loss

for i in range(20000):
    loss = train_step(X)
    if i % 2000 == 0:
        # 查看训练后的混合权重，用softmax转成合法概率
        current_weights = tf.nn.softmax(weight_logits).numpy()
        print(f"Step {i}, Loss: {loss.numpy():.4f}, Weights: {current_weights}")

补充说明 #

如果一定要使用probs参数传递权重，需要手动加softmax约束保证概率合法性，修改Categorical定义部分即可：

# 定义无约束的权重变量
raw_weights = tf.Variable([0.2, 0.8], name='raw_weights')
# 每次调用时通过softmax转成合法概率分布
cat = tfd.Categorical(probs=tf.nn.softmax(raw_weights))

这种方式和使用logits参数效果完全等价，只是显式做了softmax转换。

你之前两种写法的错误原因：

• 传入[tf.Variable(0.2, name='weight1'),tf.Variable(0.8, name='weight2')]：列表形式的独立变量没有被正确注册到分布的可训练变量列表，GradientTape无法追踪loss和这两个变量的关联，因此报无梯度错误
• 嵌套定义tf.Variable：内层的两个Variable会被当作静态初始值，不会加入计算图；外层的weights变量没有概率约束，更新后会出现非法值，同时梯度传递被阻断，因此无法更新

内容的提问来源于stack exchange，提问作者FunctionallyFirst