Repository navigation
refactor: Compute the sklearn Ridge decoder in factorized form - #7
KenyaOtsuka wants to merge 2 commits into
Conversation
b3a795d to
8c47da0
Compare
BenchmarksI measured a complete For the shipped
The difference becomes much larger for convolutional features. At full For a scaled At full DeepRecon scale, the feature/model/output I/O implied by the production The benchmark also verifies that both implementations produce matching decoded Note: these benchmarks were run on the original factorized implementation preserved in bench/pr7-benchmarks. The subsequent consolidation of the factorized decoder into a single model.pkl.gz artifact does not change the factorization or feature-loading strategy measured here. |
8d2c06b to
cfdc71e
Compare
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, info.yaml
<decoder>/<layer>/<subject>/<roi>/model/ y_mean.mat, y_norm.mat
model.pkl.gz is the decoder, whole: the fitted estimator, the ordered
training labels its coefficient columns are indexed by, the format name
and version, and the provenance of the fit, in one pickle. bdpy's
ModelTest reads only 'model' and 'y_shape' from such a file and ignores
the rest, so the extra keys ride along at no cost -- and there is no
second file that could disagree with the coefficients.
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
The ordered labels sit in the same pickle as the coefficients they
index, so the two cannot be reordered, renamed or replaced against each
other. What is left is the assembly of F: TrainingFeatureLoader.get
requires 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 check_column_correspondence compares the widths 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. A model.pkl.gz on the shared path is the
discriminator -- the direct implementation stores its model per layer,
never there -- and the format tag inside it is checked when it is
loaded; 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 and no change to evaluation.py or cv_evaluation.py. The
baseline the factorized output is checked against is the reference
implementation in tests/helpers/legacy_ridge.py, which
tests/test_legacy_reference.py pins against the pre-factorization
scripts. The benchmarks that produced the numbers quoted in the pull
request are kept outside the repository.
"Decoding with generic regression models" was never generic: the only training script in that section is train_decoder_sklearn_ridge.py, and after the factorization it is Ridge by construction, right above a paragraph that says so. The heading now names the estimator, like the PyFastL2LiR section above it. The cross-validation section called cv_train_decoder_fastl2lir.py an example for scikit-learn Ridge regression; it imports and fits FastL2LiR. That parenthetical and the equivalent one on the sklearn Ridge trainer are dropped, matching the PyFastL2LiR bullets, which carry none. predict_feature.py is left described as it was: it is not Ridge-specific -- its legacy path predicts any estimator stored in bdpy's pickle layout -- so only the claim about the pipeline is corrected.
cfdc71e to
884f330
Compare
HirokiYasuda03
left a comment
HirokiYasuda03
left a comment
There was a problem hiding this comment.
Thank you!
I left some comment.
| 'decoder.features.paths of the config used for training).' | ||
| % model_dir) | ||
|
|
||
| makedir_ifnot(results_dir_prediction) |
There was a problem hiding this comment.
Although this is the problem from previous version, in the factorized path, the layer's training features are loaded after the output directory is created and the lock is taken. If the load fails (OOM on a large layer, a missing label, a wrong feature index), the script exits but leaves an empty decoded_features/<layer>/<sbj>/<roi>/ behind. The next run then treats it as "already done" and skips it, so the gap only shows up at evaluation.
Could we either load the training features before makedir_ifnot, or remove the directory (and release the lock) on failure?
| % (label, row.shape[0], layer)) | ||
| rows.append(row) | ||
|
|
||
| features = np.vstack(rows) |
There was a problem hiding this comment.
Peak memory here is ~2x the layer size: the per-label rows list and the stacked array exist at the same time. When the statistics are missing, np.std adds another full-size temporary, so the peak is ~3x. For large layers (~15 GB per layer per the docstring) that is ~45 GB at prediction time, which the legacy predictor never needed.
Suggestions:
- Preallocate the output array and fill it row by row.
- Compute mean/std in unit (column) blocks, similar to the blocked approach in bdpy#145.
That would keep the peak close to 1x the layer size.
| feat_labels = np.unique(brain_labels) | ||
| feat = get_multi_features(data_features, layer, labels=feat_labels) |
There was a problem hiding this comment.
Training no longer touches the feature store, so a training label without features (e.g. a blank/catch trial) is only detected at prediction time. By then the decoder is marked done and has to be retrained.
The old code had no explicit check either: it failed by accident with an IndexError from feat_index. The new loader's error is much clearer. Could we run the same label-only existence check here, so the failure happens before training rather than after?
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
Repeated fMRI trials are still used individually; only the target representation
changes.
What
(subject, ROI).M × n_voxelsmodel instead of coefficients that scale with the DNNfeature dimension.
and fit metadata together in a single
model.pkl.gzartifact.coefficients and the training features.
prediction path.
evaluation.pyand the decoded-feature output format are unchanged.Compatibility / behavior changes
Test
101 tests pass on Python 3.10 and 3.11, 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, soincluding the intercept. The existing feature normalization also cancels
exactly. The full derivation is in
ridge_factorization.py.The factorized model is stored once at
while the existing per-layer
y_mean.mat/y_norm.matfiles remain underfor
evaluation.py.The fitted estimator and the ordered training-stimulus labels are stored
together in
model.pkl.gz. At prediction time, the training features are loadedin that recorded order, and their row count is checked against the predicted
coefficient dimension before they are combined.
Legacy decoders are detected automatically and use the previous prediction
path.