feat: opt-in gradient-safe log-det via Settings (default unchanged)#392
Merged
Conversation
…lt unchanged) The two Bayesian-evidence log-det terms (log_det_curvature_reg_matrix_term and log_det_regularization_matrix_term) use 2*sum(log(diag(cholesky(M)))), which returns NaN on the JAX backend (raises on numpy) when M is not numerically positive-definite -- stalling gradient-based searches on ill-conditioned regularization matrices (autolens_workspace_developer#104). Add Settings.log_det_method (Optional -> conf general.yaml fallback, mirroring the use_positive_only_solver pattern): - "cholesky" (default) -- the historical computation, byte-for-byte unchanged including the test-mode LinAlgError guard. Both terms route through a shared _log_det_symmetric_from helper; the default branch is the identical expression. - "slogdet" -- logabsdet of xp.linalg.slogdet(M): identical where M is positive-definite, finite and differentiable where the Cholesky would NaN. Opt-in, non-default, general (serves the adaptive schemes too); intended for gradient-based work and for comparison against the Cholesky evidence. Scope decided by an independent adversarial review of the reg-logdet investigation: the current evidence's coefficient-dependence is correct to machine precision, so the DEFAULT is deliberately not changed. This adds an opt-in alternative only. Full test_autoarray suite 926 passed (every existing log_evidence test unchanged = byte-identical default). New tests: default==cholesky==numpy log det; slogdet==cholesky on a PD matrix; slogdet finite where cholesky fails on a non-PD matrix. No workspace config mirroring needed -- packaged-only keys layer through production config (verified, like nnls_target_kappa). Closes #391. 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
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.
Closes #391.
Adds an opt-in, non-default gradient-safe log-determinant method to inversions, selected via
Settings. The default evidence path is unchanged and byte-identical.Why
The two Bayesian-evidence log-det terms compute
2*sum(log(diag(cholesky(M)))). On the JAX backend the Cholesky returns NaN (numpy raises) whenMis not numerically positive-definite — which stalls gradient-based searches on ill-conditioned regularization matrices (autolens_workspace_developer#104: the pixelized likelihood's NaN wall localised to exactly this term).Why opt-in and not a default change
An independent adversarial review of the reg-logdet investigation established that the current evidence is scientifically correct —
d log det(H) / d log λ = 1798to machine precision, so the coefficient-dependence that drives inference is right, and the1e-8conditioning term is only a constant offset. A relative lift would be a regression; a pseudo-determinant would break archived-evidence comparability. So the default is deliberately preserved and this adds an alternative only, for gradient work and for comparison.Change
Settings.log_det_method, following the existingOptional[…] = None → conf general.yamlpattern (mirrorsuse_positive_only_solver):"cholesky"(default) — the historical computation, byte-for-byte unchanged including the test-modeLinAlgErrorguard. Both terms now route through a shared_log_det_symmetric_fromhelper; the default branch is the identical expression."slogdet"—logabsdetofxp.linalg.slogdet(M): identical whereMis positive-definite, finite and differentiable where the Cholesky would NaN. General (serves the adaptive schemes, unlike an analytic precomputed-eigenvalue route). Opt-in viaal.SettingsInversion(log_det_method="slogdet").Applies to both evidence log-det terms (
log_det_curvature_reg_matrix_termandlog_det_regularization_matrix_term) — a gradient search needs both finite.API Changes
Additive only. New optional
Settings(log_det_method=...)field, defaultNone→ config"cholesky"→ current behaviour. No removed/renamed/changed symbols; no downstream migration.Detail
Settings.log_det_method(optional kwarg + property);AbstractInversion._log_det_symmetric_from(private helper);inversion.log_det_method: choleskyin packagedgeneral.yaml.log_det_method="slogdet"opt-in differs.nnls_target_kappa).Testing
test_autoarray/suite: 926 passed. Every existinglog_evidencetest (fit_imaging / fit_interferometer / fit_dataset) unchanged — this is the byte-identical-default gate.log_det_method == "cholesky"and reproducesnumpy log det;slogdet == choleskyon a positive-definite matrix (~1e-8);slogdetfinite (== np.linalg.slogdet) on a non-PD matrix where Cholesky fails.slogdetmatches the truelog detmore closely than Cholesky (~2.5e-7 on a 25-pixel Constant reg matrix, slogdet being the accurate one) — a further reason it is non-default and for comparison rather than a silent swap.Follow-ups (not in this PR)
autogalaxy_workspace_test(mirrorsautofit_workspace_test/scripts/jax_assertions/fitness_nan_gradient_contract.py).ConstantZerothis dead code (shape + missing-arg), andAdapt's double-square 4th-power coefficient.Gate
PyAutoLens: 1 uncommitted source change(s)(a straypaper_jax/paper.mdedit in the main checkout, untouched by this branch).