Drop the masked logsumexp reduction in mixture log_prob - #2250
Conversation
The where= kwarg of jax.nn.logsumexp in _MixtureBase.log_prob forces a masked reduction that defeats XLA's fused reduction path (3x slower for MixtureSameFamily log_prob at 10k values on CPU). The mask is redundant: the value-preserving jnp.where introduced in pyro-ppl#1874 already yields the same values and stops gradients into -inf component log-probs; the where= kwarg was added incidentally in the type-hints PR pyro-ppl#2032. Values and gradients are unchanged, including at zero mixing weights and for all-components -inf rows; add regression tests for both cases.
|
Thanks @kyo219! Can you share the code used for benchmarks. |
Benchmark reportthis PR + run time: 1 faster
compile time: unchanged across 32 benchmarksSignificant changes (1) ──────── run time ─────── ────── compile time ─────
benchmark baseline this PR Δ baseline this PR Δ
────────────────────────────────────────────────────────────────────────────────────────
+ mixture_same_family_log_prob 2.1 ms 981 µs -52.6% 102.6 ms 81.3 ms -20.8%Red is slower, green is faster; a row is coloured by the worse of its two columns. A delta in parentheses cleared the threshold on a measurement below the resolution floor, so it is shown without being called a change. † marks a benchmark that could not be compared — see below. Full results
|
| baseline | this PR | |
|---|---|---|
| ref | master |
mixture-logprob-unmasked-logsumexp |
| commit | 999d8d1f |
91d46a6e |
| numpyro | 0.21.0 | 0.21.0 |
| jax | 0.11.1 | 0.11.1 |
| backend | cpu | cpu |
| python | 3.14.7 | 3.14.7 |
Runner: Linux-6.17.0-1022-azure-x86_64-with-glibc2.39, 4 CPUs.
Produced by this benchmark run.
|
Sure! Benchmark script below (Apple M2, CPU, jax 0.10.2; jitted steady-state, best of 300 calls after warmup). The same setup was used to check the value/gradient agreement mentioned in the PR description (log_prob values, grads, zero-mixing-weight grads, and all-components import time
import jax, jax.numpy as jnp
import numpyro.distributions as dist
N = 10_000
mix = dist.Categorical(probs=jnp.array([0.2, 0.3, 0.1, 0.25, 0.15]))
comp = dist.Normal(jnp.arange(5.0), jnp.ones(5))
m = dist.MixtureSameFamily(mix, comp)
v = jax.random.normal(jax.random.PRNGKey(0), (N,)) * 3
lp = jax.jit(m.log_prob)
g = jax.jit(jax.grad(lambda v: m.log_prob(v).sum()))
def bench(f, x, n=300):
jax.block_until_ready(f(x))
t0 = time.perf_counter()
for _ in range(n):
r = f(x)
jax.block_until_ready(r)
return (time.perf_counter() - t0) / n
print(f"log_prob: {bench(lp, v)*1e3:.3f} ms grad: {bench(g, v)*1e3:.3f} ms")On the CI benchmark report: the |
2 participants
Changes made
numpyro/distributions/mixtures.py:_MixtureBase.log_probpassedwhere=~jnp.isneginf(sum_log_probs)tojax.nn.logsumexp. A masked reduction prevents XLA from using its fused reduction path, which makes the reduction the dominant cost of mixture likelihood evaluation. The mask is redundant: the value-preservingjnp.where(jnp.isneginf(...), -inf, ...)introduced by gh-1870: Refactorlog_probmethod in_MixtureBaseclass to handle-jnp.inf#1874 already produces the same values and, throughwhere's select-based gradient, stops gradients from flowing into-infcomponent log-probs. Thewhere=kwarg was not part of that fix — it slipped in with the type-hints PR feat(gh-299): Type hints in distributions modules #2032. This change restores the gh-1870: Refactorlog_probmethod in_MixtureBaseclass to handle-jnp.inf#1874 formulation and adds a comment explaining why the value-preservingwheremust stay.Values and gradients are unchanged (verified identical for regular values, at zero mixing weights, and for rows where every component is
-inf).Benchmarks
Apple M2 (24 GB), CPU, jax 0.10.2, Python 3.11. Jitted steady-state per-call time, best of 300 calls after warmup, 10,000 values, 5 components (3 for
MixtureGeneralwith StudentT).MixtureSameFamily(Categorical, Normal)log_probMixtureGeneral(5 × Normal) log_probMixtureGeneral(2 × Normal + StudentT) log_probThe big win is
MixtureSameFamily, where the component log-probs come from one broadcast producer and the masked reduction forced it to be rematerialized per pass. The last row shows a small absolute regression (~0.02 ms) for one composition where the component log_prob dominates; it appears to be an XLA fusion-choice artifact rather than a systematic cost.Links to related issues/PRs
#1874 introduced the
-infgradient protection this PR preserves; #2032 incidentally added thewhere=kwarg this PR removes.Tests
test_mixture_log_prob_grad_at_zero_weightsandtest_mixture_log_prob_grad_all_components_neg_infintest/test_distributions_mixture.py— lock in the gh-1870: Refactorlog_probmethod in_MixtureBaseclass to handle-jnp.inf#1874 semantics (no NaN gradients at zero mixing weights and for all--infrows). Neither case was previously covered by a test.pytest test/test_distributions_mixture.py— 47 passedruff check/ruff format --checkclean;ty checkreports the same 10 pre-existing diagnostics as master.Dependencies
None.