Monte Carlo Gradient Estimators

Estimate the gradient of an expectation in two ways, with torch.distributions.

Let z ~ Normal(mu, sigma) and let f be a function. You want the gradient of E[f(z)] with respect to mu and sigma. There are two standard Monte Carlo estimators:

pathwise:        z = mu + sigma * eps,  eps ~ Normal(0, 1)
                 gradient of mean(f(z))

score function:  z ~ Normal(mu, sigma), held fixed
                 gradient of mean(f(z) * log p(z; mu, sigma))

The pathwise estimator moves the samples with the parameters, so it needs f to be differentiable. The score function estimator moves only the probability of each sample, so it works for any f, but it is much noisier.

Write both. Make mu and sigma float64 tensors that require a gradient, build a torch.distributions.Normal, and draw n_samples samples. Use rsample for the pathwise estimator and sample for the score function estimator. Return the two gradient estimates as Python floats.

Function signature

def pathwise_grad(f, mu, sigma, n_samples):
def score_function_grad(f, mu, sigma, n_samples):

Example

Input

torch.manual_seed(0)
pathwise_grad(lambda z: z ** 2, mu=1.5, sigma=0.8, n_samples=200_000)

Output

(3.00..., 1.60...)    # close to (3.0, 1.6)

Reasoning

E[z ** 2] = mu ** 2 + sigma ** 2, so the exact gradient is (2 * mu, 2 * sigma) = (3.0, 1.6). The estimate is within a few thousandths.

Constraints and notes

• mu and sigma > 0 are Python floats. f takes a float64 tensor of samples and returns a tensor of the same shape.
• The tests set the seed before each call. Return (d_mu, d_sigma) as floats.
• A baseline in the score function estimator is allowed, but not needed.