Bugfix/null mixture-weights grad#2221
Closed
esennesh wants to merge 3 commits into
Closed
Conversation
Introduce a shared `_CategoricalBase` for `CategoricalProbs` and `CategoricalLogits`, holding the common `mean`/`variance`/`support`/ `enumerate_support` keyed off a `_param` hook (the raw stored parameter, so shape/dtype queries never trigger a probs<->logits conversion). Each subclass exposes `probs` without routing through the log-probs: `CategoricalProbs` returns the stored weights directly (exact zeros preserved), `CategoricalLogits` returns `softmax(logits)`. This gives a mixing distribution a `probs` property that is safe to differentiate at exact-zero weights, which the mixture log_prob fix relies on. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
_MixtureBase.log_prob wrote the density in log-weight form via log_softmax(mixing.logits) + log p_k. For a probs-parameterized mixing Categorical with an exact-zero weight, logits takes log(0): the forward value was masked to finite but the VJP evaluated 1/0 = inf, producing a NaN gradient into every parameter feeding the weights. Compute log_prob as a numerically stable weighted logsumexp instead, log p(x) = m + log sum_k w_k exp(log p_k(x) - m), reading mixing.probs (linear weights, never log(w_k)). This is exact in both value and gradient: d/dw_k = p_k(x) / sum_j w_j p_j(x) stays finite at w_k == 0, and d/d(log p_k) is the responsibility. The direct sum is used rather than jax.nn.logsumexp's `b=`, which would silently zero the gradient of any exactly-zero-weight component. Out-of-support components (log p_k = -inf) drop out with an exact zero, and an all-out-of-support value yields -inf with a zero (not NaN) gradient. A private _bare_component_log_probs returns log p_k(x); component_log_probs keeps its documented contract, now shared in _MixtureBase as log(mixing.probs) + bare. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
A mixture must evaluate every component's density at the same value, so the component support must be identical across components and independent of their parameters. The previous guard required the support itself to be a ParameterFreeConstraint, which rejected event-wrapped components such as Normal(...).to_event(1) whose support is Independent(Real(), 1) -- even though an Independent wrapper only reinterprets batch dims as event dims and introduces no parameter dependence. Unwrap any Independent layers and require the *base* support to be parameter-free instead. This enables per-element vector mixtures (e.g. a per-pixel RGB mixture whose assignment is shared across channels) with no other changes: mixture_dim, _bare_component_log_probs, component_sample and mean are already event-aware. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
So basically, there's a bugfix and a mild improvement/feature here:
probsinCategorical, even when the probabilities were specified in linear space and one of them could be a hard zero. This was causing some NaN issues.Independentdistributions when they're given as mixture components.