Fix degenerate MSE loss on regression datasets (issue #42) (#45)

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.

Author: PierreGtch

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

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)

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)