Skip to content

Implement dual ridge regression for memory optimization#2364

Open
KartikP wants to merge 5 commits into
masterfrom
kp/dual-ridge
Open

Implement dual ridge regression for memory optimization#2364
KartikP wants to merge 5 commits into
masterfrom
kp/dual-ridge

Conversation

@KartikP
Copy link
Copy Markdown
Collaborator

@KartikP KartikP commented Apr 10, 2026

Ridge and RidgeCV regression now use the dual kernel form when n_samples < n_features. This is common in Brain-Score which there are only hundreds of stimuli but thousands to hundreds of thousands of model features. This avoids materializing the coefficient matrix (n_features, n_targets) which can produce OOM for large models on fMRI benchmarks.

This PR adds DualRidgeRegression and DualRidgeCVRegression in the metrics/regression_correlation/metric.py.

Instead of solving for the full weight matrix, we precompute a small project matrix and apply it to the neural target in chunks.

This PR effectively reduces a coefficient matrix for ViT-L at 148k features on Allen2022 at 52k neuroids stored at 8 bytes from 62GB down to ~1MB.

Validation

  • All 4 Allen2022 surface ridge benchmarks produce identical raw scores
  • Identical predictions across alpha values
  • RidgeCV selects the sample alpha as sklearn

@KartikP KartikP requested a review from KogniJannis April 15, 2026 16:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant