Skip to content

refactor: Compute the sklearn Ridge decoder in factorized form - #7

Draft
KenyaOtsuka wants to merge 2 commits into
KamitaniLab:devfrom
KenyaOtsuka:refactor/factorized-sklearn-ridge-decoder
Draft

refactor: Compute the sklearn Ridge decoder in factorized form#7
KenyaOtsuka wants to merge 2 commits into
KamitaniLab:devfrom
KenyaOtsuka:refactor/factorized-sklearn-ridge-decoder

Conversation

@KenyaOtsuka

@KenyaOtsuka KenyaOtsuka commented Sep 9, 2026 •
edited
Loading

Copy link
Copy Markdown

This PR is intended to be merged after #6.

Why

The sklearn Ridge decoder stores coefficients that scale with the DNN feature
dimension, making DeepRecon-scale decoders very large and expensive to train.

This PR factorizes the same Ridge estimator as

brain → coefficients over training stimuli → decoded feature

Repeated fMRI trials are still used individually; only the target representation
changes.

What

  • Fit and store one Ridge model per (subject, ROI).
  • Store an M × n_voxels model instead of coefficients that scale with the DNN
    feature dimension.
  • Reconstruct features at prediction time from the predicted stimulus
    coefficients and the training features.
  • Training no longer reads DNN features.
  • Previously trained sklearn Ridge decoders remain usable through the legacy
    prediction path.
  • evaluation.py and the decoded-feature output format are unchanged.

Compatibility / behavior changes

  • Factorized prediction requires the training feature directory.
  • Constant feature units now decode instead of failing on zero variance.
  • The training Python API no longer takes feature-side arguments.

Test

108 tests pass, including direct-vs-factorized equivalence, regression against
the previous implementation, legacy decoder compatibility, and the unchanged
evaluation pipeline.

Benchmark results and scripts: https://github.com/KenyaOtsuka/feature-decoding/tree/bench/pr7-benchmarks

Reviewer details

For training targets Y = R F, Ridge is affine in the target, so

Ridge(alpha).fit(X, R @ F).predict(X_test)
==
Ridge(alpha).fit(X, R).predict(X_test) @ F

including the intercept. The existing feature normalization also cancels
exactly. The full derivation is in ridge_factorization.py.

The factorized model is stored once at

<decoder>/<subject>/<roi>/model/

while the existing per-layer y_mean.mat / y_norm.mat files remain under

<decoder>/<layer>/<subject>/<roi>/model/

for evaluation.py.

Prediction verifies the correspondence between stored training-stimulus labels,
loaded training features, and coefficient columns before combining them.

Legacy decoders are detected automatically and use the previous prediction path.

@KenyaOtsuka
test: Add pytest infrastructure and regression tests for the sklearn …
6938115 
…Ridge decoder

The repository had no automated tests, which made it unsafe to change the
decoding implementation. This adds the first test suite, pinning the current
behavior of `train_decoder_sklearn_ridge.py` and `predict_feature.py`.

- Add `scikit-learn` to the project dependencies. It is imported by
  `train_decoder_sklearn_ridge.py` but was never declared, so `uv sync`
  produced an environment that could not run the script.
- Add a `dev` dependency group with pytest and `[tool.pytest.ini_options]`,
  including `pythonpath = ["."]` so the top-level analysis modules are
  importable from the tests without any `sys.path` manipulation.
- `tests/helpers/synthetic.py`: build tiny `BData` / `Features` datasets on
  disk, with repeated stimuli (unbalanced repeats), two ROIs, a subject split
  over two `.h5` files, and both a 2-D `fc`-like and a 3-D `conv`-like feature
  layer so the chunking / `order='F'` path is exercised.
- `tests/helpers/legacy_ridge.py`: a plain numpy + scikit-learn reference
  implementation of the direct brain -> feature Ridge decoder, mirroring bdpy's
  `ModelTraining` / `ModelTest` semantics exactly.
- `tests/generate_golden.py` and `tests/data/golden/`: regression fixtures
  recording the pipeline output on a fixed-seed synthetic dataset. Run it as
  `python -m tests.generate_golden`.
- Tests covering the artifact contract, decoded-feature output format,
  normalization parameters, label alignment, chunking, and the behavior of
  feature units with zero training variance.
- Add a GitHub Actions workflow running the suite on push and pull requests.

No change to the decoding implementation.
@KenyaOtsuka
KenyaOtsuka force-pushed the refactor/factorized-sklearn-ridge-decoder branch 3 times, most recently from 9f895ee to b3a795d Compare September 10, 2026 01:58
@KenyaOtsuka @claude
refactor: Factorize the sklearn Ridge decoder
8c47da0 
The decoder regressed brain activity directly onto DNN features and
pickled the fitted Ridge, so the stored coefficient matrix was
d_out x n_voxels: ~120 GB for one VGG19 conv1_1 decoder of one
(subject, ROI) at DeepRecon scale, and one Ridge solve per layer and per
feature chunk.

None of that is intrinsic. Each stimulus is presented in several trials,
so the training target is not one independent vector per trial but the
features of the M unique stimuli, repeated: Y = R F with R the one-hot
trial-to-stimulus assignment. Ridge prediction is linear in the target
and R has unit row sums, so

  Ridge(a).fit(X, R F).predict(Xt) == Ridge(a).fit(X, R).predict(Xt) @ F

exactly, including the intercept, and the feature mean/SD round trip
cancels rather than being dropped. The full derivation is in the module
docstring of ridge_factorization.py; tests/test_ridge_factorization.py
pins it in float64 at rtol=1e-9 across alphas, repetition patterns,
solvers and chunk axes.

So the decoder is trained and stored as brain -> training stimulus
basis, and prediction combines the coefficients with F. The stored model
is M x n_voxels, it no longer depends on the feature layer, and the
repeated trials are *not* averaged: all N trials enter the fit exactly
as before, only the target is re-expressed.

Artifact layout. The model is the same for every layer, so it is stored
once per (subject, ROI); the training feature statistics do depend on
the layer, and evaluation.py reads them from where it always has:

  <decoder>/<subject>/<roi>/model/          model.pkl.gz, x_mean.mat,
                                            x_norm.mat, train_labels.json,
                                            factorized.yaml, info.yaml
  <decoder>/<layer>/<subject>/<roi>/model/  y_mean.mat, y_norm.mat

y_mean/y_norm are compatibility artifacts for the evaluation pipeline,
not parameters of the model -- prediction never reads them, because the
normalization cancels. Training cannot produce them (it opens no feature
file), so predict_feature.py writes them on its first run and never
overwrites what is there: a decoder directory carries the statistics of
one training feature configuration, and another feature set or feature
index is another decoder directory.

The fit is on R, built from the trial labels, so training takes no
feature-side argument at all: no feature paths, no feature index, no
chunk_axis (which only blocks the coefficient-feature product in memory
at prediction), not even the layer names. Over a DeepRecon run that
removes ~3 TB of reads that only ever recomputed identical statistics.

What the artifact does pin is the correspondence the coefficients need:

  coefficient column i is multiplied by the feature row of stimulus i

enforced by a digest of the ordered train_labels.json, by
TrainingFeatureLoader.get requiring exactly one feature row from exactly
one store per requested label (bdpy's get_multi_features appends a row
for every store holding a label and skips labels no store holds, so a
duplicate and a missing label cancel out in the total row count while
every row after them is shifted), and by check_column_correspondence at
the contraction.

Prediction resolves the decoder by looking at the shared directory
first: a factorized decoder's per-layer directory holds nothing but the
statistics sidecars, so from the outside it looks exactly like a legacy
decoder directory. The presence of factorized.yaml on the shared path is
the only discriminator; decoders written by the direct implementation
keep their layout and their previous code path, unchanged.

Deliberate behavior change: a feature unit that is constant across the
training stimuli made the direct implementation fail (0/0 = NaN targets
that ModelTraining's inf clean-up does not catch, which scikit-learn
rejects) and now decodes exactly, since C·1 = 1. That is a bug fix, not
a consequence of the refactor, and tests/test_zero_variance_features.py
pins both behaviors side by side.

Also fixes feature_index_file, which could not run at all: the feature
stores resolved the index inside the feature directory while the copy
into the decoder resolved it against the current directory. It is now
prediction-side unit selection, resolved once through
resolve_feature_index_path and copied next to the decoded features. The
same unresolved-path bug remains in train_decoder_fastl2lir.py,
predict_feature_fastl2lir.py, evaluation.py and the two cv_* scripts,
left alone here.

No bdpy changes, no change to evaluation.py or cv_evaluation.py, and the
golden fixtures from the direct implementation are reproduced unchanged. The
benchmarks that produced the numbers quoted in the pull request are kept
outside the repository.
@KenyaOtsuka
KenyaOtsuka force-pushed the refactor/factorized-sklearn-ridge-decoder branch from b3a795d to 8c47da0 Compare September 10, 2026 05:35
@KenyaOtsuka

KenyaOtsuka commented Sep 10, 2026

Copy link
Copy Markdown
Author

Benchmarks

I measured a complete train → predict run for the legacy and factorized
decoders. Training and prediction were run in separate processes, with a cold
page cache. Full scripts, methodology, and results are available on
bench/pr7-benchmarks.

For the shipped fc6 / fc7 / fc8 configuration (1200 training stimuli ×
5 repetitions, 10,000 voxels, 2 subjects × 2 ROIs):

legacy factorized change
training time 248.5 s 57.9 s 4.3× faster
prediction time 10.3 s 11.3 s about the same
total time 258.8 s 69.2 s 3.7× faster
peak memory (training) 2.9 GB 3.0 GB about the same
peak memory (prediction) 362 MB 290 MB 1.2× lower
total bytes read 2.6 GB 1.6 GB 1.7× less
total bytes written 1.4 GB 194 MB 7.3× less
decoder size 1.4 GB 184 MB 7.7× smaller

The difference becomes much larger for convolutional features. At full
conv5_1 size (100,352 units, 6000 trials, 10,000 voxels), legacy training was
OOM-killed on a 15.7 GB machine. The factorized pipeline completed on the
same machine: training took 14.0 s with 2.5 GB peak RSS, prediction took
12.5 s with 1.1 GB peak RSS, and the decoder was 45.9 MB.

For a scaled conv5_1 run where both variants fit in memory, total time was
51.0 → 9.6 s (5.3× faster); prediction itself was also slightly faster
(6.6 → 5.2 s), while training peak memory fell from 3.7 GB to 781 MB.

At full DeepRecon scale, the feature/model/output I/O implied by the production
access pattern is 51.7 TB for the legacy implementation vs 236 GB for the
factorized one (arithmetic from array sizes, not a timing measurement).

The benchmark also verifies that both implementations produce matching decoded
features (rtol=1e-4) and matching normalization statistics.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@KenyaOtsuka