Skip to content

← JAX Stochasticity step 20 of 25

Hard Primitives

Importance Sampling

Jump to the task

Why this matters

Importance sampling estimates an expectation under a target distribution p using samples drawn from a proposal distribution q. Each sample's contribution is multiplied by p(z) / q(z): a point that is relatively more likely under the target receives more weight.

For normalized densities and adequate proposal support, E_p[h(z)] = E_q[h(z) * p(z) / q(z)]. The proposal must assign positive density wherever the target's contribution to this expectation is nonzero. Reweighting cannot recover a region that the proposal never samples.

This exercise only performs the final arithmetic. The log densities and function values have already been evaluated at the same sample points.

Worked mini-example

Suppose two sample points have weights 2 and 0.5, and function values 3 and 4. Their weighted contributions are 6 and 2. Average those two contributions to get an estimate of 4.

import jax.numpy as jnp

log_q_at = jnp.log(jnp.array([0.1, 0.2]))
log_p_at = jnp.log(jnp.array([0.2, 0.1]))
h_at = jnp.array([3.0, 4.0])
weights = jnp.exp(log_p_at - log_q_at)  # Approximately [2, 0.5].
estimate = jnp.mean(h_at * weights)    # Approximately 4.0.

Dividing the sum of contributions by the sum of weights would instead give 8 / 2.5 = 3.2. That is the self-normalized estimator, a different calculation. The two estimates can differ even when both densities are normalized.

Problem

Implement importance_estimator(log_p_at, log_q_at, h_at) to return:

jnp.mean(h_at * jnp.exp(log_p_at - log_q_at))

All three inputs are finite float32 vectors of the same nonzero length N. Entries at index i refer to the same sample point. Function values may be positive, negative or zero. Assume the weights, products and reduction remain representable in float32. Return a float32 scalar.

Do not draw new samples or normalize the weights to sum to one. The fixtures check the supplied arrays' weighted mean; those arrays alone cannot establish that the underlying densities are normalized or that the proposal has adequate support.

Common pitfalls

  • Divide by the sample count N, not by the sum of weights. Ordinary importance sampling uses normalized target and proposal densities. If the target is only known up to a constant, its ordinary estimate carries that unknown scale; self-normalization cancels it but is generally biased at finite sample sizes.
  • Subtract log densities before exponentiating. This avoids separately underflowing tiny density values, but exp(log_p_at - log_q_at) can still overflow if the difference is too large.
  • Uneven weights can let a few samples dominate. A proposal with lighter tails than the target can even give an estimator with infinite variance. Equal weights of e at all supplied points merely rescale their average; they do not demonstrate a few samples dominating.
  • Density values at selected points do not have to sum to one. They are not necessarily probabilities over an exhaustive list of discrete states.
  • When corresponding log densities are equal, every weight is one and the result reduces to mean(h_at).