Full-covariance Gaussian for simplex-constrained variables - #114
Open
alexeid wants to merge 1 commit into
Open
Conversation
The default event space bijector for a Dirichlet (e.g. nucleotide base frequencies) is SoftmaxCentered, which appends a fixed zero coordinate before the softmax. The final component of the constrained variable therefore carries no unconstrained coordinate of its own and is determined by the others. A factorised (mean field) Gaussian in the unconstrained space is not invariant to which component is the pivot, and it misestimates the dispersion of that component -- typically underestimating it substantially. Under the conventional ACGT ordering the pivot is T, which is why the frequency of T is conspicuously under-dispersed relative to MCMC while A, C and G are approximated well. This gives simplex-constrained variables a Normal block with a full (rather than diagonal) covariance in their unconstrained space, leaving the remaining variables factorised. Because changing the pivot is a linear reparameterisation of the unconstrained coordinates and the full-covariance Normal family is closed under linear maps, the approximation becomes invariant to the choice of pivot and the marginal uncertainties are recovered. The cost is negligible: the block has dimension one less than the number of components (3x3 for four frequencies). Simplex variables are detected structurally (any variable whose event space bijector contains a SoftmaxCentered), so this needs no model-specific configuration. It is a hybrid by default -- full covariance only for the simplex blocks, mean field elsewhere -- exposed via the new `full_covariance_names` argument: "simplex" (default), a set of names for manual control, or None for the previous fully factorised behaviour. Verified against a Dirichlet target with closed-form Beta marginals (test/model/approximation/test_simplex_approximation.py): mean field under-disperses the pivot (~0.6x the true sd) while full covariance recovers every marginal to within ~15% and is invariant to component permutation. On the paper's datasets the VI/MCMC sd ratio for T moves from ~0.60 to ~0.93 (carnivores) and ~0.52 to ~0.84 (H3N2), matching the other frequencies.
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.
What
Gives simplex-constrained variables (e.g. nucleotide base frequencies) a Normal block with full (rather than diagonal) covariance in their unconstrained space, leaving all other variables factorised. Detected structurally — any variable whose event-space bijector contains a
SoftmaxCentered. Exposed via a newfull_covariance_namesargument toget_mean_field_approximation/get_fixed_topology_mean_field_approximation:"simplex"(default), an explicit set of names, orNonefor the previous fully-factorised behaviour.Why
The default event-space bijector for a Dirichlet is
SoftmaxCentered, which appends a fixed zero coordinate before the softmax. The final component of the constrained variable therefore has no unconstrained coordinate of its own and is determined by the others. A factorised Gaussian in the unconstrained space is not invariant to which component is the pivot, and it underestimates that component's dispersion. Under the conventional ACGT ordering the pivot is T — which is exactly why the frequency of T is conspicuously under-dispersed relative to MCMC while A, C and G are fine.Because changing the pivot is a linear reparameterisation of the unconstrained coordinates and the full-covariance Normal family is closed under linear maps, the full-covariance block makes the approximation invariant to the pivot and recovers the marginal uncertainties. Cost is negligible: the block is one dimension smaller than the number of components (3×3 for four frequencies).
Verification
test/model/approximation/test_simplex_approximation.pyfits a Dirichlet target with closed-form Beta marginals: mean field under-disperses the pivot (~0.6× true sd) while full covariance recovers every marginal to within ~15% and is invariant to component permutation. On the paper datasets the VI/MCMC sd ratio for T moves from ~0.60→0.93 (carnivores) and ~0.52→0.84 (H3N2), matching the other frequencies.Relationship to #108
This is the hybrid option discussed by email — full covariance only on the simplex blocks, mean field elsewhere — and is complementary to the
full_rank.py(full covariance over all variables) added in #108. Different files, no conflict; the two can coexist and you can choose which to make the default.