如何高效计算Score Function梯度估计器（REINFORCE算法）

需求与原理 #

我需要使用Score Function梯度估计器（REINFORCE算法），该方法无需直接对函数f求导，而是通过蒙特卡洛近似，结合带参数θ的分布p的对数概率密度梯度，来估计函数f期望的梯度。

期望梯度示意图：

近似解示意图：

性能瓶颈问题 #

在使用TensorFlow自动微分和TensorFlow Probability（TFP）分布实现该估计器时，遇到了严重的性能瓶颈。以下是示例代码（分布p为标准正态分布，f(x)=x为恒等函数）：

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

# 分布设置
scale = 1
mean = 0

normal_dist = scipy.stats.norm(loc=mean, scale=scale)

# 独立数据点数量
T = 1000

seed=360
rng = np.random.RandomState(seed)
y_T = normal_dist.rvs(size=T, random_state=rng)

# 创建Keras模型
inputs = keras.Input(shape=0)
# 无输入的全连接层仅作为截距项
intercept_layer = keras.layers.Dense(1, activation = None)
# 创建均值等于输入的正态分布
distribution_layer = tfp.layers.DistributionLambda(lambda t: tfp.distributions.Normal(loc=t, scale=1))
predicted_mean = intercept_layer(inputs)
distribution = distribution_layer(predicted_mean)
model = keras.Model(inputs=inputs, outputs=distribution)

# 创建空输入以执行前向传播
x_TF = np.zeros(shape=(T,0))

# Score Function技巧的采样数量
N=50
with tf.GradientTape() as tape:
    model_distribution = model(x_TF)
    sample_y_NT = model_distribution.sample(N)
    sample_log_probs_NT = model_distribution.log_prob(sample_y_NT)

# 使用雅可比矩阵获取每个样本的梯度
jacobian_params_NT = tape.jacobian(sample_log_probs_NT , model.trainable_weights)

上述代码在N=25时约耗时40秒，N=50则无法运行完成；尝试使用@tf.function装饰器后性能反而更差。由于标准正态分布的对数概率密度函数解析求导非常简单，本应支持更大的样本量，因此问题应该出在GradientTape的使用方式上。

我的环境配置为：Keras 2.x、TensorFlow 2.15、TensorFlow Probability 0.23，同时也接受JAX或PyTorch的解决方案。

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