Repository navigation
CHS: per-observation scores by AD, for BHHH and ML inference - #107
Open
hmgaudecker wants to merge 3 commits into
Open
hmgaudecker wants to merge 3 commits into
hmgaudecker wants to merge 3 commits into
Conversation
`scores(params)` returns the derivatives of the per-observation log-likelihoods (observations x parameters). Observations are independent, so the gradient of one observation's log-likelihood, vmapped over observations, gives all scores at about the cost of one gradient. The data are split as for the gradient. BHHH needs the scores at every iteration; so far they were only available by finite differences. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq
For an optimizer that works on likelihood contributions (BHHH), `estimate_chs` passes the per-observation scores, marked as a likelihood Jacobian, instead of the scalar log-likelihood and its gradient. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq
estimate_ml computed the score matrix of the OPG covariance by finite
differences, about one likelihood evaluation per free parameter (an hour
on a real model). The vmapped AD scores cost about one gradient.
estimate_ml_options={"jacobian": None} restores finite differences.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## interval-measurements #107 +/- ##
=========================================================
+ Coverage 96.62% 96.65% +0.02%
=========================================================
Files 161 161
Lines 18122 18182 +60
=========================================================
+ Hits 17511 17573 +62
+ Misses 611 609 -2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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.
Stacked on #106.
What
get_maximization_inputsreturnsscores: the derivatives of the per-observation log-likelihoods (observations × parameters). It isjax.gradof one observation's log-likelihood, vmapped over observations, so it costs about one gradient. The data slices (split_dataset) are shared withloglike_and_gradient.estimate_chspasses the scores toestimate_mlasjacobian, so the OPG covariance no longer uses finite differences.estimate_ml_options={"jacobian": None}restores them.jac; all others keepfun_and_jac.Why
On the health-cognition five-capacities translog (17,776 persons, 1,986 free parameters), estimagic's finite-difference Jacobian took about an hour after every fit. The AD scores take 5.4 s on an A100, at 16.3 GiB peak.
Tests
test_likelihood_regression.py).estimate_chsgives the scores to inference and to likelihood optimizers, and the override works (test_chs_estimate.py).https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq