Flux.jl是否有PyTorch中rsample等价实现支持重参数化可微分采样？

Flux.jl生态可以通过Distributions.jl + DistributionsAD.jl + Zygote实现与PyTorch rsample完全等价的功能——即基于重参数化技巧的可微随机采样，自动支持策略梯度计算。

核心原理 #

重参数化技巧的本质是将随机采样分解为无梯度的标准分布采样 + 可导的参数变换。例如正态分布的采样可表示为：
z = μ + σ * ε，其中ε ~ N(0,1)（固定采样，无梯度），μ和σ是可训练参数，变换过程完全可导，因此能通过自动微分框架计算参数梯度。

具体实现 #

1. 依赖安装 #

首先安装所需包：

using Pkg
Pkg.add(["Flux", "Distributions", "DistributionsAD", "Zygote"])

2. 手动实现重参数化采样（以正态分布为例） #

如果你想手动控制重参数化逻辑，可以定义如下函数：

using Flux, Distributions, Zygote

function rsample(d::Normal)
    # 从标准正态分布采样无梯度的ε
    ε = rand(Normal(0, 1))
    # 重参数化变换，返回可微的采样结果
    return d.μ + d.σ * ε
end

# 测试梯度计算
μ = Flux.param(2.0)
σ = Flux.param(exp(1.0))  # 用exp确保σ为正
d = Normal(μ, σ)

sample = rsample(d)
loss = sum(sample .^ 2)
grads = Zygote.gradient(() -> loss, μ, σ)

println("μ的梯度: ", grads[1])
println("σ的梯度: ", grads[2])

3. 使用DistributionsAD简化实现 #

DistributionsAD是Distributions的自动微分扩展，已经内置了多种分布的重参数化可微采样，直接调用rand即可等价于rsample：

using DistributionsAD

# 使用ADNormal（支持自动微分的正态分布类型）
d = ADNormal(μ, σ)
sample = rand(d)  # 等价于PyTorch的rsample

# 同样可以计算梯度
loss = sum(sample .^ 2)
grads = Zygote.gradient(() -> loss, μ, σ)

支持的分布类型 #

DistributionsAD支持绝大多数常见分布的重参数化可微采样，包括但不限于：

• 正态分布（ADNormal）
• 对数正态分布（ADLogNormal）
• 均匀分布（ADUniform）
• 伽马分布（ADGamma）

注意事项 #

• 必须使用DistributionsAD提供的分布类型（前缀AD），或手动实现重参数化逻辑，否则Zygote无法正确计算梯度。
• 采样过程中从标准分布生成的ε不需要求导，Zygote会自动忽略这部分的梯度传递。

内容的提问来源于stack exchange，提问作者Jose Manuel de Frutos