Struct ReparamGradEstimator 

pub struct ReparamGradEstimator {
    pub config: ReparamGradConfig,
    pub baseline: f64,
    pub n_estimates: usize,
    pub grad_variance: Vec<f64>,
}

Reparameterization gradient estimator for ELBO optimization

This struct provides a clean interface for estimating the gradient of the ELBO using the reparameterization trick. It supports:

1. Vanilla reparameterization: z = f(epsilon; phi), grad_phi ELBO via chain rule
2. Variance reduction with control variates: Uses an exponential moving average baseline to reduce gradient variance
3. Rao-Blackwellization: Analytically marginalizes when possible

Fields

config: ReparamGradConfig

Configuration

baseline: f64

Baseline value for control variate (exponential moving average of ELBO)

n_estimates: usize

Total number of gradient estimates computed

grad_variance: Vec<f64>

Running estimate of gradient variance (per-parameter)

Implementations

impl ReparamGradEstimator

pub fn new(config: ReparamGradConfig) -> Self

Create a new reparameterization gradient estimator

pub fn estimate<F>( &mut self, family: &dyn VariationalFamily, log_joint_fn: F, step: usize, ) -> Result<(f64, Array1<f64>)>

where F: Fn(&Array1<f64>) -> Result<(f64, Array1<f64>)>,

Estimate the ELBO gradient w.r.t. variational parameters phi.

Uses the reparameterization trick: grad_phi E_q[f(z)] = E_{epsilon}[grad_phi f(g(epsilon; phi))] where z = g(epsilon; phi) is the reparameterization.

Arguments

• family: The variational family q(z; phi)
• log_joint_fn: Closure returning (log p(x,z), grad_z log p(x,z))
• step: Current optimization step (for seed variation)

Returns

(elbo_estimate, grad_phi)

pub fn gradient_snr(&self) -> Vec<f64>

Compute the signal-to-noise ratio (SNR) of the gradient estimate.

High SNR (> 1) indicates stable gradient estimates. Low SNR suggests increasing n_samples or using variance reduction.

Trait Implementations

impl Clone for ReparamGradEstimator

fn clone(&self) -> ReparamGradEstimator

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for ReparamGradEstimator

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.