如何对TensorFlow Distributions进行参数化 将Categorical设为Normal的loc参数

注意不能直接把Categorical分布实例传给Normal的loc参数，loc需要接收数值型张量，你需要的是将Categorical的采样输出作为loc的取值，实现给离散类别结果加高斯噪声的需求，具体实现方式如下：

首先导入基础依赖：

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions

场景1：固定类别对应固定均值（作为Normal的loc源） #

如果每个离散类别对应固定的均值数值，可按以下两种方式实现：

• 方式一：显式分步采样

# 定义3个类别的概率
cat_probs = [0.2, 0.5, 0.3]
categorical_dist = tfd.Categorical(probs=cat_probs)

# 每个类别对应的均值，对应Normal的loc参数
class_means = tf.constant([1.0, 5.0, 10.0])
# 高斯噪声的标准差，可按需调整
noise_std = 0.5

# 单步采样流程
def sample_noisy_categorical():
    # 先采样得到类别索引
    cat_sample = categorical_dist.sample()
    # 取出对应类别的均值作为Normal的loc
    selected_loc = tf.gather(class_means, cat_sample)
    # 从带指定loc的Normal中采样，得到加了高斯噪声的结果
    noisy_sample = tfd.Normal(loc=selected_loc, scale=noise_std).sample()
    return noisy_sample

• 方式二：用Mixture分布构造端到端可调用的分布实例（支持直接采样、计算对数概率，更适合训练场景）

noisy_categorical_dist = tfd.Mixture(
    cat=categorical_dist,
    components=[
        tfd.Normal(loc=mu, scale=noise_std) 
        for mu in class_means
    ]
)

# 直接调用API即可完成对应操作
samples = noisy_categorical_dist.sample(1000) # 采样1000个样本
log_probs = noisy_categorical_dist.log_prob(samples) # 计算样本的对数概率

场景2：可训练的类别均值 #

如果需要在模型训练中更新类别对应的均值（也就是Normal的loc参数），将对应参数设为tf.Variable即可：

# 可训练的类别均值，会在反向传播时自动更新
trainable_class_means = tf.Variable([1.0, 5.0, 10.0])
# 可训练的噪声标准差，用Softplus约束为正数
noise_std = tfp.util.TransformedVariable(0.5, tfp.bijectors.Softplus(), name="noise_scale")
# 也可以用logits参数定义可训练的类别概率
trainable_cat = tfd.Categorical(logits=tf.Variable([0.0, 0.0, 0.0]))

# 构造可训练的带噪分布
trainable_noisy_dist = tfd.Mixture(
    cat=trainable_cat,
    components=[
        tfd.Normal(loc=mu, scale=noise_std)
        for mu in trainable_class_means
    ]
)

如果你是要给Categorical输出的one-hot编码加高斯噪声，而不是给类别对应的数值加噪声，只需要把class_means换成one-hot矩阵即可：

class_means = tf.eye(3) # 3个类别对应3维one-hot向量

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