Differentiating through simulations / reference request

[statsybanksy]

I have been a Jax user for a long time now, but only recently did I realise that it allows differentiating through monte carlo simulations.

Having thought about how that's possible & then googled, I am struggling to find intuitive, engineering or mathematical explanations of how that actually works. [For background, I do understand the classic automatic differentiation.]

I was wondering if you have any intuition or resources you could share on the topic. I appreciate I probably won't understand (or don't even need to understand) the nitty-gritty details of how it works, but would be especially keen to understand where the limitations of the approach are.

To make it specific, below is a simple example of a monte carlo simulation in jax I've coded up.

import jax
import jax.numpy as jnp
from jax import random, grad

# Monte Carlo simulation to estimate E(adj_sin(X))
def monte_carlo_sin(mu, sigma, num_samples=10000, key=random.PRNGKey(0)):
    samples = mu + sigma * random.normal(key, (num_samples,))
    sin_samples = jnp.sin(samples)
    adj_sin_samples = jnp.where(sin_samples > 0, 0.05 * sin_samples, sin_samples)
    return jnp.mean(adj_sin_samples)

# Example run
mu = 0.0
sigma = 1.0
key = random.PRNGKey(42)

exp = monte_carlo_sin(mu, sigma)
grad_mu = grad(monte_carlo_sin)(mu, sigma)

print(f"Exp: {exp}")
print(f"dE(adj_sin(X))/dmu: {grad_mu}")

[jakevdp]

Hi - thanks for the question!

I'm not sure exactly what answer you're looking for here – differentiation through a simulation is no different than any other autodiff setting. You're computing the partial derivative with respect to mu, and so all other inputs (including the random key) are held fixed. So for example this line:

samples = mu + sigma * random.normal(key, (num_samples,))

is essentially this with respect to autodiff:

samples = mu + const

and the autodiff rule for __add__ specifies that d(samples)/d(mu) = 1.

Other expressions in the function are handled similarly, and then the chain rule is used to propagate the gradients through the sequence of operations to get the autodiff result.

Does that make sense?