[PERF] Vectorize the likelihood estimators across parallel datasets - #137
Open
jdkent wants to merge 3 commits into
Open
[PERF] Vectorize the likelihood estimators across parallel datasets#137jdkent wants to merge 3 commits into
jdkent wants to merge 3 commits into
Conversation
Both likelihood estimators looped over the second dimension of y in Python
and ran scipy.optimize.minimize once per dataset, which put them 200-3000x
behind the closed-form estimators: ~9 s per 1000 datasets for the
variance-based one and ~135-205 s for the sample-size-based one, against
~0.03 s for WLS, DL and Hedges.
Both likelihoods are now profiled down to a single bounded parameter and
searched for every dataset at once:
- The coefficients have a closed form at a fixed variance component, so
substituting it leaves a function of tau^2 alone. That is exactly
equivalent to minimizing the joint likelihood.
- The sample-size-based likelihood is invariant to the scale of
(sigma^2, tau^2) up to a closed-form factor, so only their ratio, which
lies in [0, 1], has to be searched for. Three parameters become one.
- tau^2 is searched through u / (1 - u), a monotone map from [0, 1) onto
[0, inf), so the bounded search cannot truncate the parameter space. The
D-L moment estimate sets the scale of that map, which is the vectorized
counterpart of warm-starting a per-dataset optimizer.
- stats.bounded_scalar_min does a coarse scan to bracket each dataset's
minimum and then refines every bracket in step by golden section, so the
whole fit costs ~80 vectorized evaluations regardless of dataset count.
Measured at 30 observations, 12 groups, seconds per 1000 datasets:
before after
ungrouped resc. ungrouped resc.
VBL ML 8.93 11.75 0.54 0.54
VBL REML 8.39 10.38 0.57 0.51
SSBL ML 134.10 135.21 0.51 0.49
SSBL REML 204.92 131.54 0.56 0.53
Scored with the original joint objective, the new solution is never worse
(largest difference 1e-14) and on 6 of 200 datasets the sample-size-based
fit lands a substantially lower nll -- up to 1.59 -- than the per-dataset
optimizer reached, which is where its parameter estimates change.
The warning about looping over more than ten datasets goes away with the
loop, as does the wrapt dependency it was built on.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… them
Profiling the vectorized fits showed one call dominating every estimator in
the library: np.linalg.pinv, which decomposes each of the D copies of X'WX
by SVD. It was 48-58% of a likelihood fit and 65% of a DerSimonian-Laird
fit. Nothing else came close, and the Python-level overhead of the search is
7.5 ms per fit however large D is, so it was never dispatch-bound.
Three changes, in descending order of payoff:
- stats._invert_stack replaces the pseudo-inverse with a reciprocal when
there is one predictor and an LU inverse otherwise, falling back to
np.linalg.pinv where that is the only defined answer. Measured on a stack
of 20000 matrices: 74.7 ms for pinv, 7.1 ms for inv, 0.03 ms for 1 / x.
This one also speeds up every closed-form estimator.
- The objective no longer recomputes what the fit it just ran already has.
The weights are formed once and handed to weighted_least_squares instead
of being rebuilt inside it, and REML reads log|X'WX| off the returned
covariance as -log|(X'WX)^-1| rather than forming X'WX a second time.
- bounded_scalar_min refines by successive parabolic interpolation with
Brent's safeguards rather than by golden section, stops where the
objective has gone flat to double precision, and leaves alone the
datasets whose optimum the scan already placed at an end of the interval
-- the very common tau^2 = 0. The variance-based fit went from 78
objective evaluations to 49. Brent's step-size safeguard is load-bearing:
without it the interpolation creeps by ever-smaller steps on a nearly flat
minimum and stops converging.
Seconds per 1000 datasets, 30 observations, 12 groups, D = 2000:
before vectorized now vs vectorized vs before
VBL ML 1.707 0.147 0.0344 4.3x 50x
VBL REML 2.728 0.201 0.0404 5.0x 68x
SSBL ML 19.429 0.192 0.0634 3.0x 306x
SSBL REML 22.413 0.200 0.0770 2.6x 291x
WLS 0.0017 0.0017 0.0007 2.4x 2.4x
Hedges 0.0037 0.0031 0.0006 5.2x 6.2x
DL 0.0036 0.0028 0.0019 1.5x 1.9x
Scored with the original joint objective the fits are unchanged: never
worse than the per-dataset optimizer anywhere (largest difference 2e-13)
and better on the same 6 of 200 sample-size-based datasets as before. The
refinement was also checked against a 20001-point grid on 400 random
non-quadratic objectives, where it never returned a higher value.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #137 +/- ##
==========================================
+ Coverage 92.18% 92.35% +0.17%
==========================================
Files 13 13
Lines 1817 1845 +28
==========================================
+ Hits 1675 1704 +29
+ Misses 142 141 -1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull request overview
Vectorizes likelihood estimators across parallel datasets and accelerates matrix inversion.
Changes:
- Adds a vectorized bounded scalar minimizer.
- Profiles likelihood estimation to one parameter.
- Updates tests, benchmarks, documentation, and dependencies.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
setup.cfg |
Removes the unused wrapt dependency. |
pymare/stats.py |
Adds vectorized minimization and faster inversion. |
pymare/estimators/estimators.py |
Vectorizes likelihood estimators. |
pymare/tests/test_stats.py |
Tests minimization and inversion behavior. |
pymare/tests/test_estimators.py |
Tests parallel likelihood fitting. |
docs/api.rst |
Exposes the new minimizer API. |
benchmarks/common.py |
Removes loop-specific dataset sizing. |
benchmarks/bench_estimators.py |
Benchmarks vectorized likelihood fits. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
All three reproduce, and each now has a test that fails without its fix. The sample size identifiability guard was reducing over the whole K x D array once the per-dataset loop went away, so it asked whether sample sizes vary *anywhere* rather than within each column. Columns that are each constant still differ from one another, so a set of datasets that are every one of them unidentifiable passed the guard and came back with numbers: two columns of constant n = 20 and n = 50 fitted to tau^2 = 1.91 and 1e-4. A single constant column also slipped through beside a varying one. The old loop raised for both. Reduced over observations now, and the message says how many datasets are affected. Same for the near-equal warning. bounded_scalar_min's flat test compared the two ends of the bracket with each other rather than with the middle, which its own docstring says it does. An objective steeper on one side of its minimum than the other can hold both ends at equal height while the middle sits far below them; the search read that as flat, stopped before its first refinement step and returned the scan point. On a V with slopes 10 and 1, placed so the ends come out equal, it returned 1/3 against a true minimum of 79/264 -- out by 3.4e-2, most of a scan cell, after 31 evaluations. Now 8.6e-9 after 81. Costs nothing on real data: the evaluation counts and timings for all four likelihood fits are unchanged. The documented ``lower <= upper`` contract was not enforced. A reversed interval descends, so the tolerance changes sign with it, the first convergence test passes, and a scan point is returned as though it had been refined -- 0.333 rather than 0.314 on a quadratic. It raises now. An interval of one point is ordered, and still allowed. Co-Authored-By: Claude Opus 5 (1M context) <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.
Summary
The two likelihood estimators looped over the second dimension of
yin Python and ranscipy.optimize.minimizeonce per dataset, putting them 200–3000× behind the closed-form estimators. Measured at 30 observations and 12 groups, seconds per 1000 datasets: ~9 forVarianceBasedLikelihoodEstimatorand ~135–205 forSampleSizeBasedLikelihoodEstimator, against ~0.03 for WLS, DL and Hedges. Extrapolated to a 200k-voxel image that is roughly half an hour and eleven hours respectively.Both are now profiled down to a single bounded parameter and searched for every dataset at once, and the inversion that dominated every estimator in the library has been replaced.
What changed
Profiling the likelihoods down to one parameter.
(sigma^2, tau^2)up to a closed-form factor (Q/Kfor ML,Q/(K-P)for REML), so only their ratio, which lies in[0, 1], has to be searched for. Three parameters become one.u / (1 - u), a monotone map from[0, 1)onto[0, inf), so a bounded search cannot truncate the parameter space. The DerSimonian-Laird moment estimate sets the scale of that map, which is the vectorized counterpart of warm-starting a per-dataset optimizer.stats.bounded_scalar_minscans coarsely to bracket each dataset's minimum, then refines every bracket in step by successive parabolic interpolation with Brent's safeguards. The whole fit costs a few dozen vectorized evaluations regardless of dataset count.Not taking an SVD per objective evaluation. Profiling showed one call dominating every estimator:
np.linalg.pinv, which decomposes each of the D copies ofX'WXby SVD — 48–58% of a likelihood fit and 65% of a DerSimonian-Laird fit.stats._invert_stackuses a reciprocal when there is one predictor and an LU inverse otherwise, keepingpinvas the fallback where it is the only defined answer. On a stack of 20000 matrices: 74.7 ms forpinv, 7.1 ms forinv, 0.03 ms for1 / x.Not recomputing what the fit already has. The weights are formed once and handed to
weighted_least_squaresinstead of being rebuilt inside it, and REML readslog|X'WX|off the returned covariance as-log|(X'WX)^-1|rather than formingX'WXa second time.Timings
Seconds per 1000 datasets, 30 observations, 12 groups, D = 2000:
A 200k-voxel image now takes 8–15 s for the likelihood estimators.
Correctness
tau27.7649 for ML, 10.9499 for REML) still hold atatol=1e-4.sigma2andtau2estimates change: the old three-parameter L-BFGS-B run was under-converging.sigma^2andtau^2.Full suite: 331 passed, 1 skipped.
flake8clean.API notes
Three deliberate changes worth flagging:
**kwargson the two estimators now goes tostats.bounded_scalar_min(xtol,ftol,maxiter) rather than toscipy.optimize.minimize. Code passing SciPy minimizer options will raiseTypeErrorat fit time.wraptis dropped frominstall_requires; it existed only for the_loopabledecorator. An ill-conditioned but nonsingular design now gets the ordinary inverse rather than one with its smallest singular values truncated, so a nearly collinear design reports a large covariance instead of a quietly regularized one.The benchmark suite's likelihood entries move to the shared dataset size, since they no longer need a smaller second dimension, and
N_DATASETS_LOOPEDis removed.🤖 Generated with Claude Code