From 3c7d8e19b9316d614cc94afd850e2007e2409ce5 Mon Sep 17 00:00:00 2001 From: Pierre Guetschel Date: Fri, 26 Jun 2026 15:38:34 +0200 Subject: [PATCH] Fix degenerate MSE loss on regression datasets (issue #42) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Windowed regression targets are scalar-per-sample, so the collated batch target has shape (B,) while the regression model output is (B, 1). The default MSELoss silently broadcast (B, 1) against (B,) to (B, B), averaging over B^2 pairwise differences instead of the B element-wise ones. That broadcast loss equals the true element-wise MSE plus 2*cov(pred, target), i.e. it actively penalises prediction/target correlation and wrecks regression training, while the reported r2 metric (which ravels predictions) stayed correct — so the bug was silent. Give the SGD regression path an EEGRegressor subclass whose get_loss aligns the target's trailing dimension to the (B, 1) output. Classification (EEGClassifier) and ridge probing (already reshapes targets) are unaffected. Add tests/test_regression_loss.py guarding against the broadcast. --- CHANGELOG.md | 3 ++ open_eeg_bench/training.py | 36 +++++++++++++++++-- tests/test_regression_loss.py | 65 +++++++++++++++++++++++++++++++++++ 3 files changed, 101 insertions(+), 3 deletions(-) create mode 100644 tests/test_regression_loss.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d3013fe..f015e51 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed +- Fix degenerate MSE loss on regression datasets (e.g. `seed_vig`): the scalar-per-window target has shape `(B,)` while the model output is `(B, 1)`, so `MSELoss` silently broadcast `(B, 1)` against `(B,)` to `(B, B)` and trained against `2 * cov(pred, target)` added to the true error — actively penalising prediction/target correlation. The SGD regression path (`Training`) now aligns the target's trailing dimension to the output before the loss. Classification and ridge probing were unaffected ([#42](https://github.com/braindecode/OpenEEGBench/issues/42)). + ## [0.6.0] - 2026-06-02 ### Added diff --git a/open_eeg_bench/training.py b/open_eeg_bench/training.py index 0c60bc5..f1739a1 100644 --- a/open_eeg_bench/training.py +++ b/open_eeg_bench/training.py @@ -72,6 +72,37 @@ def _build_callback(self): return WandbLogger(wandb_run=wandb_run, save_model=False) +def _aligned_target_eeg_regressor(): + """Build an ``EEGRegressor`` subclass that aligns the target's trailing + dimension to the model output before computing the loss. + + Windowed regression targets are scalar-per-sample, so the collated batch + target has shape ``(B,)`` while the regression model output is ``(B, 1)``. + The default ``MSELoss`` would otherwise *broadcast* ``(B, 1)`` against + ``(B,)`` to ``(B, B)`` and average over ``B**2`` pairwise differences + instead of the ``B`` element-wise ones — a silent, degenerate objective. + That broadcast loss equals the true element-wise MSE *plus* + ``2 * cov(pred, target)``, i.e. it actively penalises prediction/target + correlation, which wrecks regression training. See issue #42. + + The subclass is built lazily so importing this module stays cheap (no + eager ``torch`` / ``braindecode`` import), consistent with the rest of the + package. + """ + from braindecode import EEGRegressor + from skorch.utils import to_tensor + + class _AlignedTargetEEGRegressor(EEGRegressor): + def get_loss(self, y_pred, y_true, *args, **kwargs): + y_true = to_tensor(y_true, device=self.device) + # Scalar regression target (B,) -> (B, 1) to match output (B, 1). + if y_pred.dim() == 2 and y_pred.size(1) == 1 and y_true.dim() == 1: + y_true = y_true.unsqueeze(1) + return super().get_loss(y_pred, y_true, *args, **kwargs) + + return _AlignedTargetEEGRegressor + + class Training(BaseModel): """Training hyper-parameters and callback settings.""" @@ -154,7 +185,7 @@ def build_learner(self, model, callbacks, n_classes, val_set, verbose=1, seed: i # training configs share the same `build_learner` signature. from torch.optim import AdamW from skorch.helper import predefined_split - from braindecode import EEGClassifier, EEGRegressor + from braindecode import EEGClassifier is_regression = n_classes is None callbacks = list(callbacks) @@ -174,9 +205,8 @@ def build_learner(self, model, callbacks, n_classes, val_set, verbose=1, seed: i iterator_valid__num_workers=self.num_workers, ) - is_regression = n_classes is None if is_regression: - learner = EEGRegressor(**common_kwargs) + learner = _aligned_target_eeg_regressor()(**common_kwargs) else: classes = list(range(n_classes)) learner = EEGClassifier(classes=classes, **common_kwargs) diff --git a/tests/test_regression_loss.py b/tests/test_regression_loss.py new file mode 100644 index 0000000..02e6842 --- /dev/null +++ b/tests/test_regression_loss.py @@ -0,0 +1,65 @@ +"""Regression-target loss shape: guard against the degenerate MSE broadcast. + +See issue #42. A windowed regression target is a scalar per sample, so the +collated batch target is shape ``(B,)`` while the regression model output is +``(B, 1)``. The default ``MSELoss`` then *broadcasts* ``(B, 1)`` against +``(B,)`` to ``(B, B)`` and averages ``B**2`` pairwise differences instead of +the ``B`` element-wise ones — a silent, degenerate training objective. + +The fix gives the target a trailing dimension inside the regression learner's +loss so the comparison is element-wise. +""" + +import warnings + +import torch +import torch.nn as nn + +from open_eeg_bench.training import Training + + +def _build_regressor(): + """Build a regression learner (n_classes=None) on a tiny dummy module.""" + val_set = torch.utils.data.TensorDataset( + torch.zeros(4, 4), torch.zeros(4) + ) + learner = Training(device="cpu").build_learner( + model=nn.Linear(4, 1), + callbacks=[], + n_classes=None, + val_set=val_set, + verbose=0, + ) + learner.initialize() + return learner + + +def test_regression_loss_is_elementwise_not_broadcast(): + """get_loss on (B,1) pred vs (B,) target must be element-wise MSE.""" + learner = _build_regressor() + y_pred = torch.zeros(8, 1) + y_true = torch.arange(8, dtype=torch.float32) # shape (B,) + + with warnings.catch_warnings(): + # The broadcasting UserWarning must NOT fire. + warnings.simplefilter("error") + loss = learner.get_loss(y_pred, y_true) + + # Element-wise MSE of (0 - y_true)**2 == mean(y_true**2). + expected = (y_true**2).mean() + assert torch.isclose(loss, expected), ( + f"loss={float(loss)} (broadcast MSE over B*B pairs) " + f"!= element-wise {float(expected)}" + ) + + +def test_regression_loss_matches_2d_target(): + """A (B,) target must give the same loss as an already-(B,1) target.""" + learner = _build_regressor() + torch.manual_seed(0) + y_pred = torch.randn(8, 1) + y_true = torch.randn(8) + + loss_1d = learner.get_loss(y_pred, y_true) + loss_2d = learner.get_loss(y_pred, y_true.reshape(8, 1)) + assert torch.isclose(loss_1d, loss_2d)