perf(glm): make GLM fit faster - #442
Open
FelipeCybis wants to merge 11 commits into
Open
Conversation
Fit the runs of a multi-run FirstLevelModel through joblib.Parallel (loky workers), one run per task. Loky caps each worker's BLAS/OpenMP thread pool at cpu_count() // n_jobs, so parallel runs do not oversubscribe the CPU; n_jobs=1 keeps the sequential in-process path. h5py-backed runs cannot be pickled to workers and raise TypeError when n_jobs != 1, mirroring register_volumewise. Closes #310
Replace the batched pseudoinverse of the per-voxel normal equations with a batched inverse when the design is full rank. `pinv` runs an SVD per voxel and dominated the fit; `inv` uses LU and agrees to ~1e-14. A rank-deficient design still takes `pinv`, which is what produces the min-norm solution, and an exactly singular batch falls back to it. Contract the lagged cross-product and the three-operand XtX term with `optimize=True` so they route through BLAS instead of einsum's own loop. The remaining einsums in that block are no faster with it, and the two-operand XtX term is slower, so they keep the default. test_ar_degenerate_design only asserted degrees of freedom, so it passed regardless of the coefficients returned. Replace it with a comparison against an explicit per-voxel lstsq solve, parametrized over full-rank and collinear designs so both inversion branches are covered.
Fitting runs in joblib worker processes made multi-run fits slower rather than faster. joblib memmap-dumps every run to disk before any compute starts, and loky caps each worker's BLAS pool at cpu_count() // n_jobs, so the single large GEMM that the sequential path spreads over every core is split into narrow ones instead. Measured on 20 CPUs with four 705 MB runs, n_jobs=4 took 27.3 s against 7.6 s sequentially. Drop the parameter, its h5py pickling guard and this module's joblib dependency, and fit the runs in a plain loop. Add show_progress in its place, which is the part of the joblib path worth keeping. joblib_progress advances its bar by patching Parallel.print_progress, so it cannot tick without a Parallel to patch; add _utils/progress.py exposing the same rich columns for loops driven by hand. decoding/searchlight.py hand-rolls this same bar and can move onto it later. Feed estimate_ar_coeffs the OLS whitened residuals instead of the `residuals` property. OLS whitening is the identity, so the property recomputed `Y - design @ theta` and allocated a second (n_volumes, n_voxels) array holding values the fit had already stored.
With minimize_memory=True, which is the default, the fit built a whitened residual array of (n_volumes, n_voxels), summed its squares, and then dropped it again in _strip_heavy_fields. Since beta solves the normal equations, the same sum is `||wY||² - beta·XtY`, computed from quantities the fit already holds. Add keep_residuals to OLSModel.fit and ARModel.fit so the caller says whether the diagnostic arrays are wanted, and let RegressionResults take whitened_residuals=None. sse/mse/residuals raise in that case with the message they already raise after stripping. The first OLS pass of an AR fit always keeps its residuals, since the AR coefficients are estimated from them. Accumulate every sum of squares in float64 through _sum_of_squares. The normal-equations form is a difference of two large, close sums, so a float32 accumulation on float32 recordings loses about five significant digits; it is also faster than np.sum(a**2, axis=0), which materializes the squared array. Measured back to back on one 600 x 147k run: ARModel.fit 3.83 s -> 2.29 s and OLSModel.fit 0.88 s -> 0.59 s. Agreement with the explicit residual path is within 8e-12 relative, worst case, at R² = 0.995.
…quations Reverts 1b84795. The identity `rss = ||wY||² - beta·XtY` holds in exact arithmetic, but atlas-resampled recordings carry large constant regions outside the recorded field of view which the design's constant regressor explains completely. There the two terms are equal, so the sign of their difference is pure rounding: on the five-run example in docs/examples/06_glm, 448195 of 1203840 voxels came out with a negative residual sum of squares, hence a negative dispersion and NaN across 39% of the contrast map. It was buying 1.7 s of a 51 s fit, so the shortcut is not worth a guard. The remaining changes on this branch already take that example from 129 s to 51 s. Keep a regression test with half its voxels set to exact linear combinations of the design, scaled so the cancellation dominates the true residual. Constant voxels alone are not enough: with a single degenerate voxel the sign of the difference is a coin flip, and the first version of this test passed against the very code it was meant to reject.
`_positive_reciprocal` guarded against dividing by zero with `np.where`, which selects a branch only after both have been evaluated, so the reciprocal was still taken over the zeros the guard exists to discard. Any recording carrying voxels the design explains exactly -- an atlas grid leaves whole constant regions outside the recorded field of view -- printed a spurious divide-by-zero RuntimeWarning from every contrast. `np.divide` with `where=` skips those elements instead. The returned values are unchanged.
A progress bar on stdout is a surprise for a library call inside a script or a notebook, so the default is now False.
…d helper The column list was duplicated between searchlight and the GLM. The searchlight keeps its own comment on why the bar is advanced by hand: the inner cross_val_score builds a Parallel of its own, so joblib_progress would count its folds against this bar.
# Conflicts: # docs/changelog.md
FelipeCybis
changed the title
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
AR whitening left-multiplies the design by I - sum_i rho[i]*S^(i+1), which is unit lower triangular and therefore invertible for every rho. A full-rank design keeps every per-voxel XtX positive definite, so the batched inv cannot raise and the pinv fallback behind it was dead code. Extend the lstsq comparison with a rho far outside the stationary range, which is the case the fallback was guarding.
Member
Author
Contributor
|
📖 Doc preview: https://confusius.tools/pr-preview/pr-442/ |
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.
In my machine,
docs/examples/06_glm/01_first_level.pyfits in about 55 seconds on this branch, against about 129 seconds on main. The larger part of the gain applies to the defaultnoise_model="ar1". The z-maps are unchanged, except on voxels where rounding controls the value in both versions.Closes #310.
The three changes
1.
invin place ofpinvfor the per-voxel normal equationsARModel.fitinverts a(147k, 15, 15)stack. The pseudo-inverse is a batched SVD and costs about 6 times more than the LU inverse. The two results agree to 1e-14.pinvstays for a rank-deficient design, whereinvreturns a wrong answer without an error. The full-rank branch needs no guard of its own. Whitening multiplies the design by a unit lower triangular factor, so it preserves rank for every rho andXtXstays positive definite.Before:
After:
2. No second pass over the OLS residuals
The
residualsproperty returnsself.Y - self.predicted, which allocates a second(n_volumes, n_voxels)array. OLS whitening is the identity, sowhitened_residualsalready holds those values.Before:
After:
3.
optimize=Trueon twonumpy.einsumcallsThe lagged cross-product goes to BLAS and runs about 30 times faster. The three-operand term contracts in pairs, in place of building a full
(order, order, V, K, K)intermediate. The two-operand term between them is faster without the flag, so it keeps the default.Before:
After:
Two smaller items in this PR
FirstLevelModelacceptsshow_progress, which shows a progress bar over the runs (defaults toFalse). The bar lives inconfusius/_utils/progress.py, anddecoding/searchlight.pynow draws its own bar from that helper.compute_contrastno longer prints a divide-by-zeroRuntimeWarning. A recording that holds a voxel of no variance over time triggered this. An atlas grid leaves whole regions like that outside the recorded field of view.Before:
After: