From 6997fcb0e34a0a2c535711dd4c4539909483508c Mon Sep 17 00:00:00 2001
From: Pierre Guetschel <pierre.guetschel@donders.ru.nl>
Date: Wed, 3 Jun 2026 10:01:30 +0200
Subject: [PATCH 1/2] Add LazyGaussianRandomProjection head utility

Lazy nn.Module mirroring sklearn's GaussianRandomProjection: the input
dimension is inferred from the first forward (like LazyLinear) and the
non-trainable projection matrix is materialized into a persistent buffer.

Reuses _make_projection_matrix for bit-identical sklearn scaling and
widens its type hints (seed: int | None, device: str | torch.device) to
match the new call site.
---
 open_eeg_bench/head_utils.py  | 80 +++++++++++++++++++++++++++++++++++
 open_eeg_bench/ridge_probe.py |  6 ++-
 tests/test_head_utils.py      | 75 ++++++++++++++++++++++++++++++++
 3 files changed, 160 insertions(+), 1 deletion(-)
 create mode 100644 open_eeg_bench/head_utils.py
 create mode 100644 tests/test_head_utils.py

diff --git a/open_eeg_bench/head_utils.py b/open_eeg_bench/head_utils.py
new file mode 100644
index 0000000..c707e22
--- /dev/null
+++ b/open_eeg_bench/head_utils.py
@@ -0,0 +1,80 @@
+"""Lazy random-projection head utilities.
+
+A :class:`LazyGaussianRandomProjection` mirrors scikit-learn's
+:class:`~sklearn.random_projection.GaussianRandomProjection` but is packaged as
+a lazy :class:`torch.nn.Module`: the input dimension (``n_features``) is inferred
+from the first forward, exactly like :class:`torch.nn.LazyLinear`. The projection
+matrix is a (non-trainable) buffer rather than a learnable parameter.
+"""
+
+import warnings
+
+import torch
+import torch.nn as nn
+from torch.nn.modules.lazy import LazyModuleMixin
+from torch.nn.parameter import UninitializedBuffer
+
+from open_eeg_bench.ridge_probe import _make_projection_matrix
+
+
+class LazyGaussianRandomProjection(LazyModuleMixin, nn.Module):
+    """Gaussian random projection with a lazily-inferred input dimension.
+
+    Mirrors :class:`sklearn.random_projection.GaussianRandomProjection`: the
+    projection matrix entries are drawn from ``N(0, 1 / n_components)`` (via
+    sklearn, so the matrix is bit-identical for a given ``random_state``). The
+    source dimension is inferred from ``input.shape[-1]`` on the first forward,
+    at which point ``projection_matrix`` is materialized to shape
+    ``(n_components, n_features)``.
+
+    Args:
+        n_components: target projection dimension.
+        random_state: integer seed for the projection matrix (``None`` draws a
+            fresh matrix each time, like sklearn's default).
+        device, dtype: factory kwargs for the materialized buffer, following the
+            :class:`torch.nn.LazyLinear` convention.
+    """
+
+    projection_matrix: UninitializedBuffer
+
+    def __init__(
+        self,
+        n_components: int,
+        *,
+        random_state=None,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        self.n_components = n_components
+        self.random_state = random_state
+        self.n_features = 0
+        self.register_buffer(
+            "projection_matrix",
+            UninitializedBuffer(device=device, dtype=dtype),
+        )
+
+    def initialize_parameters(self, input) -> None:  # type: ignore[override]
+        """Infer ``n_features`` from the input and materialize the matrix."""
+        if self.has_uninitialized_params():
+            with torch.no_grad():
+                self.n_features = input.shape[-1]
+                if self.n_components > self.n_features:
+                    warnings.warn(
+                        "The number of components is higher than the number of "
+                        f"features: n_features < n_components "
+                        f"({self.n_features} < {self.n_components}). The "
+                        "dimensionality of the problem will not be reduced.",
+                    )
+                self.projection_matrix.materialize((self.n_components, self.n_features))
+                P = _make_projection_matrix(
+                    n_features=self.n_features,
+                    n_components=self.n_components,
+                    seed=self.random_state,
+                    device=self.projection_matrix.device,
+                    dtype=self.projection_matrix.dtype,
+                )
+                self.projection_matrix.copy_(P)
+
+    def forward(self, input):
+        return input @ self.projection_matrix.T
diff --git a/open_eeg_bench/ridge_probe.py b/open_eeg_bench/ridge_probe.py
index ddc15a7..b2e41a7 100644
--- a/open_eeg_bench/ridge_probe.py
+++ b/open_eeg_bench/ridge_probe.py
@@ -60,7 +60,11 @@ def _balanced_class_weights(
 
 
 def _make_projection_matrix(
-    n_features: int, n_components: int, seed: int, device: str, dtype: torch.dtype
+    n_features: int,
+    n_components: int,
+    seed: int | None,
+    device: str | torch.device,
+    dtype: torch.dtype,
 ) -> torch.Tensor:
     """Return a (n_components, n_features) Gaussian random projection matrix.
 
diff --git a/tests/test_head_utils.py b/tests/test_head_utils.py
new file mode 100644
index 0000000..b3b62f1
--- /dev/null
+++ b/tests/test_head_utils.py
@@ -0,0 +1,75 @@
+"""Tests for lazy random-projection head utilities."""
+
+import numpy as np
+import pytest
+import torch
+
+
+def test_lazy_projection_matches_sklearn_gaussian():
+    """A first forward materializes a buffer bit-identical to sklearn's
+    GaussianRandomProjection, and the projection equals its transform."""
+    from sklearn.random_projection import GaussianRandomProjection
+
+    from open_eeg_bench.head_utils import LazyGaussianRandomProjection
+
+    torch.manual_seed(0)
+    N, D, k = 50, 200, 16
+    X = torch.randn(N, D, dtype=torch.float64)
+
+    # sklearn reference: components_ and transform on the same X.
+    rp = GaussianRandomProjection(n_components=k, random_state=0)
+    Xt_ref = rp.fit_transform(X.numpy())
+
+    proj = LazyGaussianRandomProjection(
+        n_components=k, random_state=0, dtype=torch.float64
+    )
+    # Before the first forward the buffer is uninitialized.
+    assert proj.has_uninitialized_params()
+
+    Xt = proj(X)
+
+    # After the first forward the buffer is materialized.
+    assert not proj.has_uninitialized_params()
+    assert Xt.shape == (N, k)
+    assert proj.projection_matrix.shape == (k, D)
+    np.testing.assert_allclose(proj.projection_matrix.numpy(), rp.components_)
+    np.testing.assert_allclose(Xt.numpy(), Xt_ref)
+
+
+def test_warns_when_n_components_exceeds_n_features():
+    """Like sklearn, projecting up (n_components > n_features) warns rather than
+    silently producing a non-reducing projection."""
+    from open_eeg_bench.head_utils import LazyGaussianRandomProjection
+
+    proj = LazyGaussianRandomProjection(n_components=8, random_state=0)
+    X = torch.randn(4, 5)  # n_features=5 < n_components=8
+    with pytest.warns(match="higher than the number"):
+        out = proj(X)
+    assert out.shape == (4, 8)
+
+
+def test_state_dict_roundtrip_into_fresh_module():
+    """The materialized matrix is a persistent buffer: saving it and loading it
+    into a fresh (still-lazy) module materializes that module identically.
+
+    This is why the matrix is an UninitializedBuffer rather than a plain
+    attribute — it rides torch's lazy state_dict machinery like LazyLinear.
+    """
+    from open_eeg_bench.head_utils import LazyGaussianRandomProjection
+
+    X = torch.randn(3, 10, dtype=torch.float64)
+    src = LazyGaussianRandomProjection(
+        n_components=4, random_state=0, dtype=torch.float64
+    )
+    y = src(X)  # materialize
+    state = src.state_dict()
+    assert "projection_matrix" in state  # persistent buffer
+
+    # A fresh module with a different seed is still uninitialized until load.
+    dst = LazyGaussianRandomProjection(
+        n_components=4, random_state=999, dtype=torch.float64
+    )
+    assert dst.has_uninitialized_params()
+    dst.load_state_dict(state)
+    assert not dst.has_uninitialized_params()
+    torch.testing.assert_close(dst(X), y)

From 298af77ecc2e24290a0cde0792e9a090c0d39659 Mon Sep 17 00:00:00 2001
From: Pierre Guetschel <pierre.guetschel@donders.ru.nl>
Date: Fri, 26 Jun 2026 14:37:03 +0200
Subject: [PATCH 2/2] Move random projection from ridge probe into Head models

Add a `max_features: int | None = 5000` field to LinearHead, MLPHead and
FlattenHead (not OriginalHead): when set, apply() inserts a
LazyGaussianRandomProjection right after nn.Flatten(1), capping the flattened
feature dimension. head.apply() gains a `seed` argument and Experiment passes
its seed so the projection is reproducible from Experiment.seed.

LazyGaussianRandomProjection now does an identity passthrough when the input
already has <= n_components features (replacing the warn-and-expand behavior),
matching the ridge probe's old "skip when D <= max_features".

The streaming ridge probe loses all projection responsibility: max_features /
projection_seed parameters, the 'projection' result key and the predict-time
re-projection are gone, and _make_projection_matrix now lives in head_utils.
RidgeProbingTraining drops its max_features field accordingly.

Also fix _initialize_lazy_modules to materialize lazy buffers (not just lazy
parameters), so a FlattenHead projection with no trailing LazyLinear is
initialized by the dummy forward.
---
 open_eeg_bench/experiment.py  |  14 +-
 open_eeg_bench/head.py        |  89 ++++++++++---
 open_eeg_bench/head_utils.py  |  95 +++++++++-----
 open_eeg_bench/ridge_probe.py |  63 +--------
 open_eeg_bench/training.py    |  30 ++---
 tests/test_head_projection.py | 124 ++++++++++++++++++
 tests/test_head_utils.py      |  17 +--
 tests/test_ridge_probe.py     | 234 +++++++++-------------------------
 8 files changed, 357 insertions(+), 309 deletions(-)
 create mode 100644 tests/test_head_projection.py

diff --git a/open_eeg_bench/experiment.py b/open_eeg_bench/experiment.py
index 36c2f31..6ec23c4 100644
--- a/open_eeg_bench/experiment.py
+++ b/open_eeg_bench/experiment.py
@@ -189,7 +189,9 @@ def run(self) -> dict:
         # ===============================================================
         # 3. Apply head
         # ===============================================================
-        self.head.apply(model, info["n_outputs"], backbone_obj.head_module_name)
+        self.head.apply(
+            model, info["n_outputs"], backbone_obj.head_module_name, seed=self.seed
+        )
 
         # ===============================================================
         # 3.1. Initialize lazy modules with a dummy forward pass
@@ -272,13 +274,21 @@ def run(self) -> dict:
 
     @staticmethod
     def _initialize_lazy_modules(model, info: dict) -> None:
-        """Run a dummy forward pass to materialize LazyLinear modules."""
+        """Run a dummy forward pass to materialize lazy modules.
+
+        Covers both lazy parameters (e.g. ``nn.LazyLinear``) and lazy buffers
+        (e.g. a head's ``LazyGaussianRandomProjection``), which a head may add
+        without any accompanying lazy parameter.
+        """
         import torch
         import torch.nn as nn
 
         has_lazy = any(
             isinstance(p, nn.parameter.UninitializedParameter)
             for p in model.parameters()
+        ) or any(
+            isinstance(b, nn.parameter.UninitializedBuffer)
+            for b in model.buffers()
         )
         if not has_lazy:
             return
diff --git a/open_eeg_bench/head.py b/open_eeg_bench/head.py
index a99a225..bfb6697 100644
--- a/open_eeg_bench/head.py
+++ b/open_eeg_bench/head.py
@@ -2,6 +2,13 @@
 
 A Head replaces the backbone's final layer with a new one suited to the
 downstream task.  ``OriginalHead`` keeps the model's built-in head.
+
+Every head except ``OriginalHead`` flattens its input and can optionally insert
+a non-trainable Gaussian random projection right after the flatten via
+``max_features``: when the flattened dimension exceeds ``max_features`` it is
+randomly projected down to ``max_features`` (otherwise the projection is an
+identity passthrough). This caps the feature dimension fed to the downstream
+layer — in particular it replaces the ridge probe's old feature-capping step.
 """
 
 from __future__ import annotations
@@ -17,19 +24,38 @@
 log = logging.getLogger(__name__)
 
 
+def _projection_layer(max_features: int | None, seed: int | None):
+    """Return a LazyGaussianRandomProjection, or None when max_features is None."""
+    if max_features is None:
+        return None
+    from open_eeg_bench.head_utils import LazyGaussianRandomProjection
+
+    return LazyGaussianRandomProjection(max_features, random_state=seed)
+
+
 class LinearHead(BaseModel):
     """Replace the head with a single linear layer."""
 
     model_config = ConfigDict(extra="forbid")
     kind: Literal["linear"] = "linear"
+    max_features: int | None = 5000
 
-    def apply(self, model, n_outputs: int, head_module_name: str) -> None:
+    def apply(
+        self, model, n_outputs: int, head_module_name: str, seed: int | None = None
+    ) -> None:
         import torch.nn as nn
 
-        new_head = nn.Sequential(nn.Flatten(1), nn.LazyLinear(n_outputs))
-        setattr(model, head_module_name, new_head)
+        layers: list = [nn.Flatten(1)]
+        proj = _projection_layer(self.max_features, seed)
+        if proj is not None:
+            layers.append(proj)
+        layers.append(nn.LazyLinear(n_outputs))
+        setattr(model, head_module_name, nn.Sequential(*layers))
         log.info(
-            "Replaced '%s' with LinearHead (n_outputs=%d)", head_module_name, n_outputs
+            "Replaced '%s' with LinearHead (n_outputs=%d, max_features=%s)",
+            head_module_name,
+            n_outputs,
+            self.max_features,
         )
 
 
@@ -40,23 +66,32 @@ class MLPHead(BaseModel):
     kind: Literal["mlp"] = "mlp"
     hidden_dim: int = 256
     dropout: float = 0.5
+    max_features: int | None = 5000
 
-    def apply(self, model, n_outputs: int, head_module_name: str) -> None:
+    def apply(
+        self, model, n_outputs: int, head_module_name: str, seed: int | None = None
+    ) -> None:
         import torch.nn as nn
 
-        new_head = nn.Sequential(
-            nn.Flatten(1),
-            nn.LazyLinear(self.hidden_dim),
-            nn.ELU(),
-            nn.Dropout(p=self.dropout),
-            nn.Linear(self.hidden_dim, n_outputs),
+        layers: list = [nn.Flatten(1)]
+        proj = _projection_layer(self.max_features, seed)
+        if proj is not None:
+            layers.append(proj)
+        layers.extend(
+            [
+                nn.LazyLinear(self.hidden_dim),
+                nn.ELU(),
+                nn.Dropout(p=self.dropout),
+                nn.Linear(self.hidden_dim, n_outputs),
+            ]
         )
-        setattr(model, head_module_name, new_head)
+        setattr(model, head_module_name, nn.Sequential(*layers))
         log.info(
-            "Replaced '%s' with MLPHead (hidden=%d, dropout=%.2f)",
+            "Replaced '%s' with MLPHead (hidden=%d, dropout=%.2f, max_features=%s)",
             head_module_name,
             self.hidden_dim,
             self.dropout,
+            self.max_features,
         )
 
 
@@ -66,21 +101,39 @@ class OriginalHead(BaseModel):
     model_config = ConfigDict(extra="forbid")
     kind: Literal["original"] = "original"
 
-    def apply(self, model: nn.Module, n_outputs: int, head_module_name: str) -> None:
+    def apply(
+        self,
+        model: nn.Module,
+        n_outputs: int,
+        head_module_name: str,
+        seed: int | None = None,
+    ) -> None:
         log.info("Keeping original head '%s'", head_module_name)
 
 
 class FlattenHead(BaseModel):
-    """Replace the head with nn.Flatten(1). For ridge probing only."""
+    """Replace the head with a flatten (+ optional projection). Ridge probing only."""
 
     model_config = ConfigDict(extra="forbid")
     kind: Literal["flatten"] = "flatten"
+    max_features: int | None = 5000
 
-    def apply(self, model, n_outputs: int, head_module_name: str) -> None:
+    def apply(
+        self, model, n_outputs: int, head_module_name: str, seed: int | None = None
+    ) -> None:
         import torch.nn as nn
 
-        setattr(model, head_module_name, nn.Flatten(1))
-        log.info("Replaced '%s' with FlattenHead", head_module_name)
+        proj = _projection_layer(self.max_features, seed)
+        if proj is None:
+            new_head = nn.Flatten(1)
+        else:
+            new_head = nn.Sequential(nn.Flatten(1), proj)
+        setattr(model, head_module_name, new_head)
+        log.info(
+            "Replaced '%s' with FlattenHead (max_features=%s)",
+            head_module_name,
+            self.max_features,
+        )
 
 
 Head = Annotated[
diff --git a/open_eeg_bench/head_utils.py b/open_eeg_bench/head_utils.py
index c707e22..e0c2469 100644
--- a/open_eeg_bench/head_utils.py
+++ b/open_eeg_bench/head_utils.py
@@ -1,34 +1,61 @@
 """Lazy random-projection head utilities.
 
-A :class:`LazyGaussianRandomProjection` mirrors scikit-learn's
-:class:`~sklearn.random_projection.GaussianRandomProjection` but is packaged as
-a lazy :class:`torch.nn.Module`: the input dimension (``n_features``) is inferred
-from the first forward, exactly like :class:`torch.nn.LazyLinear`. The projection
+A :class:`LazyGaussianRandomProjection` caps the feature dimension at
+``n_components`` via a Gaussian random projection, packaged as a lazy
+:class:`torch.nn.Module`: the input dimension (``n_features``) is inferred from
+the first forward, exactly like :class:`torch.nn.LazyLinear`. The projection
 matrix is a (non-trainable) buffer rather than a learnable parameter.
 """
 
-import warnings
-
+import numpy as np
 import torch
 import torch.nn as nn
 from torch.nn.modules.lazy import LazyModuleMixin
 from torch.nn.parameter import UninitializedBuffer
 
-from open_eeg_bench.ridge_probe import _make_projection_matrix
+
+def _make_projection_matrix(
+    n_features: int,
+    n_components: int,
+    seed: int | None,
+    device: str | torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Return a (n_components, n_features) Gaussian random projection matrix.
+
+    Uses sklearn's ``GaussianRandomProjection._make_random_matrix`` directly so
+    that scaling matches scikit-learn exactly (entries ~ N(0, 1/n_components)).
+    The private method is called on purpose: ``fit`` requires a numpy array of
+    real data, which we don't have at this stage (only the feature dimension).
+    """
+    from sklearn.random_projection import GaussianRandomProjection
+
+    rp = GaussianRandomProjection(n_components=n_components, random_state=seed)
+    P = rp._make_random_matrix(n_components, n_features)
+    return torch.from_numpy(np.asarray(P)).to(device=device, dtype=dtype)
 
 
 class LazyGaussianRandomProjection(LazyModuleMixin, nn.Module):
-    """Gaussian random projection with a lazily-inferred input dimension.
+    """Gaussian random projection that caps the feature dimension lazily.
 
-    Mirrors :class:`sklearn.random_projection.GaussianRandomProjection`: the
-    projection matrix entries are drawn from ``N(0, 1 / n_components)`` (via
-    sklearn, so the matrix is bit-identical for a given ``random_state``). The
-    source dimension is inferred from ``input.shape[-1]`` on the first forward,
-    at which point ``projection_matrix`` is materialized to shape
-    ``(n_components, n_features)``.
+    The source dimension is inferred from ``input.shape[-1]`` on the first
+    forward (like :class:`torch.nn.LazyLinear`), at which point
+    ``projection_matrix`` is materialized:
+
+    * if ``n_features > n_components``: a ``(n_components, n_features)`` Gaussian
+      random matrix (entries ~ ``N(0, 1 / n_components)``, drawn via sklearn so
+      it is bit-identical for a given ``random_state``). ``forward`` then maps
+      ``(B, n_features) → (B, n_components)``.
+    * if ``n_features <= n_components``: an **identity** matrix, i.e. a no-op
+      passthrough that leaves the features unchanged. This mirrors the ridge
+      probe's old behavior of skipping the projection when the feature
+      dimension is already at or below the cap (no expansion, no random mixing).
+      The identity is stored in full (a ``(n_features, n_features)`` buffer) so
+      the module round-trips through ``state_dict`` like any materialized lazy
+      module; this costs ``n_features²`` elements when not reducing.
 
     Args:
-        n_components: target projection dimension.
+        n_components: maximum projection dimension (the feature cap).
         random_state: integer seed for the projection matrix (``None`` draws a
             fresh matrix each time, like sklearn's default).
         device, dtype: factory kwargs for the materialized buffer, following the
@@ -59,22 +86,30 @@ def initialize_parameters(self, input) -> None:  # type: ignore[override]
         if self.has_uninitialized_params():
             with torch.no_grad():
                 self.n_features = input.shape[-1]
-                if self.n_components > self.n_features:
-                    warnings.warn(
-                        "The number of components is higher than the number of "
-                        f"features: n_features < n_components "
-                        f"({self.n_features} < {self.n_components}). The "
-                        "dimensionality of the problem will not be reduced.",
+                if self.n_features <= self.n_components:
+                    # No reduction possible: store identity → passthrough.
+                    self.projection_matrix.materialize(
+                        (self.n_features, self.n_features)
+                    )
+                    self.projection_matrix.copy_(
+                        torch.eye(
+                            self.n_features,
+                            device=self.projection_matrix.device,
+                            dtype=self.projection_matrix.dtype,
+                        )
+                    )
+                else:
+                    self.projection_matrix.materialize(
+                        (self.n_components, self.n_features)
+                    )
+                    P = _make_projection_matrix(
+                        n_features=self.n_features,
+                        n_components=self.n_components,
+                        seed=self.random_state,
+                        device=self.projection_matrix.device,
+                        dtype=self.projection_matrix.dtype,
                     )
-                self.projection_matrix.materialize((self.n_components, self.n_features))
-                P = _make_projection_matrix(
-                    n_features=self.n_features,
-                    n_components=self.n_components,
-                    seed=self.random_state,
-                    device=self.projection_matrix.device,
-                    dtype=self.projection_matrix.dtype,
-                )
-                self.projection_matrix.copy_(P)
+                    self.projection_matrix.copy_(P)
 
     def forward(self, input):
         return input @ self.projection_matrix.T
diff --git a/open_eeg_bench/ridge_probe.py b/open_eeg_bench/ridge_probe.py
index b2e41a7..2360597 100644
--- a/open_eeg_bench/ridge_probe.py
+++ b/open_eeg_bench/ridge_probe.py
@@ -59,27 +59,6 @@ def _balanced_class_weights(
     return counts.sum() / (n_classes * counts.clamp(min=1.0))
 
 
-def _make_projection_matrix(
-    n_features: int,
-    n_components: int,
-    seed: int | None,
-    device: str | torch.device,
-    dtype: torch.dtype,
-) -> torch.Tensor:
-    """Return a (n_components, n_features) Gaussian random projection matrix.
-
-    Uses sklearn's ``GaussianRandomProjection._make_random_matrix`` directly so
-    that scaling matches scikit-learn exactly (entries ~ N(0, 1/n_components)).
-    The private method is called on purpose: ``fit`` requires a numpy array of
-    real data, which we don't have at this stage (only the feature dimension).
-    """
-    from sklearn.random_projection import GaussianRandomProjection
-
-    rp = GaussianRandomProjection(n_components=n_components, random_state=seed)
-    P = rp._make_random_matrix(n_components, n_features)
-    return torch.from_numpy(np.asarray(P)).to(device=device, dtype=dtype)
-
-
 def _fit_streaming_ridge(
     model: "nn.Module",
     train_loader: "DataLoader",
@@ -87,16 +66,14 @@ def _fit_streaming_ridge(
     n_classes: int | None,
     lambdas: list[float] | None,
     device: str,
-    max_features: int | None = None,
-    projection_seed: int = 0,
     class_weight: str | None = "balanced",
     dtype: str = "float64",
 ) -> dict:
     """Fit streaming ridge probe, select λ on val, return weights + diagnostics.
 
-    If ``max_features`` is set and the backbone emits more features than that,
-    features are projected down to ``max_features`` dimensions via a Gaussian
-    random projection (seeded by ``projection_seed``) before accumulation.
+    Features are consumed as produced by ``model`` — any dimensionality
+    reduction (e.g. a random projection) is the model/head's responsibility,
+    applied upstream during the forward pass.
 
     ``class_weight`` controls per-sample weighting in the weighted-least-squares
     fit. Only meaningful when ``n_classes`` is set (classification); silently
@@ -114,8 +91,7 @@ def _fit_streaming_ridge(
     float64.
 
     Returns dict with keys: W (D,C), bias (C,), best_lambda (float),
-    val_scores (dict λ→score), lambdas (list[float]), n_classes, n_features D
-    (after projection if applied), projection (torch.Tensor | None).
+    val_scores (dict λ→score), lambdas (list[float]), n_classes, n_features D.
     """
     model.eval()
     model.to(device)
@@ -132,27 +108,12 @@ def _fit_streaming_ridge(
 
     # ----- Pass 1: accumulate (optionally weighted) sufficient statistics on train -----
     A = B = s_h = s_h2 = s_y = None
-    projection = None  # (k, D_orig); lazily built once D_orig is known
     N = 0.0  # sum of sample weights (== n_samples when unweighted)
     with torch.no_grad():
         for batch in train_loader:
             x, y = batch[0], batch[1]
             h = model(x.to(device))
             h_acc = h.to(dtype=torch_dtype)
-            if (
-                projection is None
-                and max_features is not None
-                and h_acc.shape[1] > max_features
-            ):
-                projection = _make_projection_matrix(
-                    n_features=h_acc.shape[1],
-                    n_components=max_features,
-                    seed=projection_seed,
-                    device=device,
-                    dtype=torch_dtype,
-                )
-            if projection is not None:
-                h_acc = h_acc @ projection.T  # (B, k)
             y_dev = y.to(device)
             y_enc = _encode_targets(y_dev, n_classes)
             y_acc = y_enc.to(dtype=torch_dtype)
@@ -245,7 +206,6 @@ def _fit_streaming_ridge(
         n_classes=n_classes,
         y_bar_train=y_bar,
         device=device,
-        projection=projection,
         dtype=torch_dtype,
     )  # tensor of shape (K,)
 
@@ -273,7 +233,6 @@ def _fit_streaming_ridge(
         "lambdas": lambdas,
         "n_classes": n_classes,
         "n_features": D,
-        "projection": projection,
     }
 
 
@@ -286,7 +245,6 @@ def _streaming_val_scores(
     y_bar_train: torch.Tensor,
     device: str,
     dtype: torch.dtype,
-    projection: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Accumulate per-λ metric streaming on val. Returns (K,) scores (higher=better).
 
@@ -304,8 +262,6 @@ def _streaming_val_scores(
             for batch in val_loader:
                 x, y = batch[0], batch[1]
                 h = model(x.to(device)).to(dtype=dtype)
-                if projection is not None:
-                    h = h @ projection.T
                 y_enc = _encode_targets(y.to(device), n_classes).to(dtype=dtype)
                 preds = torch.einsum("kdc,bd->kbc", Ws, h) + biases.unsqueeze(1)
                 res = preds - y_enc.unsqueeze(0)
@@ -319,8 +275,6 @@ def _streaming_val_scores(
         for batch in val_loader:
             x, y = batch[0], batch[1]
             h = model(x.to(device)).to(dtype=dtype)
-            if projection is not None:
-                h = h @ projection.T
             y_true = y.to(device).long()
             preds = torch.einsum("kdc,bd->kbc", Ws, h) + biases.unsqueeze(1)
             y_pred = preds.argmax(dim=2)  # (K, B)
@@ -356,8 +310,6 @@ def __init__(
         device: str,
         lambdas: list[float] | None,
         val_set,
-        max_features: int | None = None,
-        projection_seed: int = 0,
         class_weight: str | None = "balanced",
         dtype: str = "float64",
         verbose: int = 1,
@@ -369,8 +321,6 @@ def __init__(
         self.device = device
         self.lambdas = lambdas
         self.val_set = val_set
-        self.max_features = max_features
-        self.projection_seed = projection_seed
         self.class_weight = class_weight
         self.dtype = dtype
         self.verbose = verbose
@@ -404,8 +354,6 @@ def fit(self, train_set, y=None):
             n_classes=self.n_classes_,
             lambdas=self.lambdas,
             device=self.device,
-            max_features=self.max_features,
-            projection_seed=self.projection_seed,
             class_weight=self.class_weight,
             dtype=self.dtype,
         )
@@ -429,7 +377,6 @@ def predict(self, test_set) -> np.ndarray:
 
         W = self._result["W"]  # (D, C)
         bias = self._result["bias"]  # (C,)
-        projection = self._result.get("projection")  # (k, D_orig) or None
         torch_dtype = _resolve_dtype(self.dtype)
         loader = DataLoader(
             test_set,
@@ -446,8 +393,6 @@ def predict(self, test_set) -> np.ndarray:
             for batch in loader:
                 x = batch[0] if isinstance(batch, (list, tuple)) else batch
                 h = self.model_(x.to(self.device)).to(dtype=torch_dtype)
-                if projection is not None:
-                    h = h @ projection.T
                 y_hat = h @ W + bias  # (B, C)
                 outs.append(y_hat.cpu().numpy())
         preds = np.concatenate(outs, axis=0)
diff --git a/open_eeg_bench/training.py b/open_eeg_bench/training.py
index 0c60bc5..7f79c7c 100644
--- a/open_eeg_bench/training.py
+++ b/open_eeg_bench/training.py
@@ -191,23 +191,13 @@ class RidgeProbingTraining(BaseModel):
     select regularization strength on val, predict on test. No gradient-based
     training, no callbacks.
 
-    Memory (peak, excluding backbone and per-batch activations)
-    -----------------------------------------------------------
-    When ``max_features`` caps the feature dimension at ``k``, peak memory
-    during the fit is dominated by:
-
-    * the projection matrix ``P`` of shape ``(k, D_orig)`` — roughly
-      ``8 · k · D_orig`` bytes (depends on the backbone; zero when
-      ``D_orig ≤ k``, no projection applied).
-    * a fixed cost of **~48 · k²** bytes independent of backbone and dataset,
-      from four persistent ``k×k`` float64 matrices (``X^T X``, centered cov,
-      correlation, eigenvectors) plus the LAPACK workspace of ``torch.linalg.eigh``.
-
-    Typical fixed-cost values (projection matrix not included):
-
-    * ``max_features = 1_000``  → ~50 MB
-    * ``max_features = 5_000``  → ~1.25 GB
-    * ``max_features = 25_000`` → ~30 GB
+    The feature dimension fed to the probe is capped by the head's
+    ``max_features`` (e.g. ``FlattenHead(max_features=k)``), which random-projects
+    the flattened features down to ``k`` upstream. Peak memory during the fit is
+    then dominated by a fixed cost of **~48 · k²** bytes (four persistent
+    ``k×k`` float64 matrices — ``X^T X``, centered cov, correlation,
+    eigenvectors — plus the LAPACK workspace of ``torch.linalg.eigh``): e.g.
+    ~50 MB at ``k = 1_000``, ~1.25 GB at ``k = 5_000``, ~30 GB at ``k = 25_000``.
     """
 
     model_config = ConfigDict(extra="forbid")
@@ -218,14 +208,14 @@ class RidgeProbingTraining(BaseModel):
     num_workers: int = 0
     device: str = "cpu"
     lambdas: list[float] | None = None  # None → use the learner's default fixed log-spaced grid
-    max_features: int | None = 5000  # if set and D > max_features, Gaussian random-project to max_features
     class_weight: Literal["balanced"] | None = "balanced"  # "balanced" → sklearn-style per-class weights; None → unweighted
     dtype: Literal["float32", "float64"] = "float64"  # "float64" recommended for precision; use "float32" only if necessary (e.g. Apple MPS)
 
     def build_learner(self, model, callbacks, n_classes, val_set, verbose=1, seed: int = 0):
         from open_eeg_bench.ridge_probe import StreamingRidgeProbeLearner
 
-        # callbacks ignored — no epochs, no early stopping
+        # callbacks and seed ignored: no epochs/early stopping, and any random
+        # projection is seeded inside the head (head.apply(seed=...)) upstream.
         return StreamingRidgeProbeLearner(
             feature_extractor=model,
             n_classes=n_classes,
@@ -234,8 +224,6 @@ def build_learner(self, model, callbacks, n_classes, val_set, verbose=1, seed: i
             device=self.device,
             lambdas=self.lambdas,
             val_set=val_set,
-            max_features=self.max_features,
-            projection_seed=seed,
             class_weight=self.class_weight,
             dtype=self.dtype,
             verbose=verbose,
diff --git a/tests/test_head_projection.py b/tests/test_head_projection.py
new file mode 100644
index 0000000..7545538
--- /dev/null
+++ b/tests/test_head_projection.py
@@ -0,0 +1,124 @@
+"""Tests for the max_features random-projection layer inserted into heads."""
+
+import pytest
+import torch
+import torch.nn as nn
+
+from open_eeg_bench.head import FlattenHead, LinearHead, MLPHead, OriginalHead
+from open_eeg_bench.head_utils import LazyGaussianRandomProjection
+
+
+def _apply(head, n_outputs=3, seed=0):
+    """Apply a head onto a bare nn.Module and return the inserted head module."""
+    model = nn.Module()
+    head.apply(model, n_outputs, "head", seed=seed)
+    return model.head
+
+
+# ----------------------------- FlattenHead ---------------------------------
+
+def test_flatten_head_max_features_reduces_dim():
+    """FlattenHead(max_features=k) projects the flattened features down to k."""
+    head_mod = _apply(FlattenHead(max_features=10))
+    x = torch.randn(4, 8, 5)  # flattened D = 40 > 10
+    out = head_mod(x)
+    assert out.shape == (4, 10)
+    assert any(isinstance(m, LazyGaussianRandomProjection) for m in head_mod)
+
+
+def test_flatten_head_none_is_bare_flatten():
+    """max_features=None → plain nn.Flatten, no projection."""
+    head_mod = _apply(FlattenHead(max_features=None))
+    assert isinstance(head_mod, nn.Flatten)
+
+
+def test_flatten_head_default_is_5000():
+    """FlattenHead defaults to max_features=5000 (preserves ridge behavior)."""
+    assert FlattenHead().max_features == 5000
+
+
+@pytest.mark.parametrize(
+    "head_factory",
+    [
+        lambda: FlattenHead(max_features=10),
+        lambda: LinearHead(max_features=10),
+        lambda: MLPHead(hidden_dim=7, max_features=10),
+    ],
+    ids=["flatten", "linear", "mlp"],
+)
+def test_head_seed_determinism(head_factory):
+    """Same seed → identical projection matrix; different seed → different.
+    Holds for every head that supports max_features."""
+
+    def proj(seed):
+        head_mod = _apply(head_factory(), seed=seed)
+        head_mod(torch.randn(4, 8, 5))  # materialize (D=40 > 10)
+        lgrp = next(m for m in head_mod if isinstance(m, LazyGaussianRandomProjection))
+        return lgrp.projection_matrix.clone()
+
+    p0a, p0b, p42 = proj(0), proj(0), proj(42)
+    torch.testing.assert_close(p0a, p0b)
+    assert not torch.allclose(p0a, p42)
+
+
+# ----------------------------- LinearHead ----------------------------------
+
+def test_linear_head_max_features_inserts_projection_before_linear():
+    head_mod = _apply(LinearHead(max_features=10), n_outputs=3)
+    x = torch.randn(4, 8, 5)  # D = 40 > 10
+    out = head_mod(x)
+    assert out.shape == (4, 3)
+    assert any(isinstance(m, LazyGaussianRandomProjection) for m in head_mod)
+
+
+def test_linear_head_none_no_projection():
+    head_mod = _apply(LinearHead(max_features=None), n_outputs=3)
+    assert not any(isinstance(m, LazyGaussianRandomProjection) for m in head_mod)
+
+
+# ----------------------------- MLPHead -------------------------------------
+
+def test_mlp_head_max_features_inserts_projection():
+    head_mod = _apply(MLPHead(hidden_dim=7, max_features=10), n_outputs=3)
+    x = torch.randn(4, 8, 5)  # D = 40 > 10
+    out = head_mod(x)
+    assert out.shape == (4, 3)
+    assert any(isinstance(m, LazyGaussianRandomProjection) for m in head_mod)
+
+
+# ----------------------------- OriginalHead --------------------------------
+
+def test_original_head_rejects_max_features():
+    """OriginalHead has no projection: max_features must be forbidden."""
+    from pydantic import ValidationError
+
+    with pytest.raises(ValidationError):
+        OriginalHead(max_features=10)
+
+
+# ------------------- lazy-buffer materialization lifecycle -----------------
+
+class _PassthroughModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.head = nn.Identity()
+
+    def forward(self, x):
+        return self.head(x)
+
+
+def test_initialize_lazy_modules_materializes_projection_buffer():
+    """Experiment._initialize_lazy_modules must materialize the head's projection
+    buffer even when no UninitializedParameter is present (FlattenHead has no
+    LazyLinear after the projection)."""
+    from open_eeg_bench.experiment import Experiment
+
+    model = _PassthroughModel()
+    FlattenHead(max_features=10).apply(model, n_outputs=3, head_module_name="head", seed=0)
+    lgrp = next(m for m in model.head if isinstance(m, LazyGaussianRandomProjection))
+    assert lgrp.has_uninitialized_params()  # not yet materialized
+
+    # dummy (1, n_chans, n_times) → flatten → 4*5 = 20 > 10 → reduces
+    Experiment._initialize_lazy_modules(model, {"n_chans": 4, "n_times": 5})
+
+    assert not lgrp.has_uninitialized_params()  # materialized by the dummy forward
diff --git a/tests/test_head_utils.py b/tests/test_head_utils.py
index b3b62f1..1fee2d7 100644
--- a/tests/test_head_utils.py
+++ b/tests/test_head_utils.py
@@ -1,7 +1,6 @@
 """Tests for lazy random-projection head utilities."""
 
 import numpy as np
-import pytest
 import torch
 
 
@@ -36,16 +35,18 @@ def test_lazy_projection_matches_sklearn_gaussian():
     np.testing.assert_allclose(Xt.numpy(), Xt_ref)
 
 
-def test_warns_when_n_components_exceeds_n_features():
-    """Like sklearn, projecting up (n_components > n_features) warns rather than
-    silently producing a non-reducing projection."""
+def test_passthrough_when_n_components_geq_n_features():
+    """When n_components >= input dim, the projection is a no-op identity
+    passthrough — matching the ridge probe's old 'skip when D <= max_features'.
+    No dimensionality expansion, no random matrix."""
     from open_eeg_bench.head_utils import LazyGaussianRandomProjection
 
     proj = LazyGaussianRandomProjection(n_components=8, random_state=0)
-    X = torch.randn(4, 5)  # n_features=5 < n_components=8
-    with pytest.warns(match="higher than the number"):
-        out = proj(X)
-    assert out.shape == (4, 8)
+    X = torch.randn(4, 5)  # n_features=5 <= n_components=8 → passthrough
+    out = proj(X)
+    assert out.shape == (4, 5)  # dimension unchanged, NOT expanded to 8
+    torch.testing.assert_close(out, X)  # exact identity passthrough
+    assert not proj.has_uninitialized_params()
 
 
 def test_state_dict_roundtrip_into_fresh_module():
diff --git a/tests/test_ridge_probe.py b/tests/test_ridge_probe.py
index ffeb07e..9d9e4db 100644
--- a/tests/test_ridge_probe.py
+++ b/tests/test_ridge_probe.py
@@ -11,11 +11,7 @@
 
 @pytest.fixture
 def classif_data():
-    """Synthetic linearly separable data for classification.
-
-    D is chosen > 100 so that the ``max_features=100`` parametrization
-    actually triggers a random projection in tests that exercise it.
-    """
+    """Synthetic linearly separable data for classification."""
     torch.manual_seed(0)
     rng = np.random.default_rng(0)
     N, D, C = 500, 200, 3
@@ -38,27 +34,21 @@ def _loader(X, y, batch_size=32):
     return DataLoader(ds, batch_size=batch_size, shuffle=False)
 
 
-@pytest.mark.parametrize("max_features", [None, 100])
-def test_fit_streaming_ridge_matches_sklearn_classification(classif_data, max_features):
-    """Streaming ridge predictions match a sklearn pipeline on the same X.
+def test_fit_streaming_ridge_matches_sklearn_classification(classif_data):
+    """Streaming ridge predictions match a sklearn StandardScaler + Ridge pipeline.
 
-    The reference is ``[GaussianRandomProjection?] → StandardScaler → Ridge``:
-    when ``max_features`` is set, the sklearn pipeline projects using the same
-    ``random_state`` as our ``projection_seed``. Identical seeds must yield
-    identical projection matrices — so fitting on raw X_tr on both sides and
-    predicting on raw X_tr must match end-to-end.
+    Random projection is no longer the probe's concern (it moved into the head),
+    so this compares the bare ridge solve against sklearn on the raw features.
     """
     from sklearn.linear_model import Ridge
     from sklearn.preprocessing import StandardScaler
     from sklearn.pipeline import make_pipeline
-    from sklearn.random_projection import GaussianRandomProjection
     from open_eeg_bench.ridge_probe import _fit_streaming_ridge
 
     X_tr, y_tr = classif_data["train"]
     X_val, y_val = classif_data["val"]
     C = classif_data["C"]
     lam = 1.0
-    projection_seed = 0
 
     out = _fit_streaming_ridge(
         model=nn.Identity(),
@@ -67,34 +57,16 @@ def test_fit_streaming_ridge_matches_sklearn_classification(classif_data, max_fe
         n_classes=C,
         lambdas=[lam],
         device="cpu",
-        max_features=max_features,
-        projection_seed=projection_seed,
         class_weight=None,  # sklearn Ridge has no class_weight; disable to compare 1-to-1
     )
 
-    if max_features is not None:
-        # make sure the test is actually testing the projection logic
-        assert max_features < classif_data["D"]
-
-    # sklearn reference pipeline: same projection (if any) + standardize + ridge.
-    steps = []
-    if max_features is not None:
-        steps.append(
-            GaussianRandomProjection(
-                n_components=max_features, random_state=projection_seed
-            )
-        )
-    steps.extend([StandardScaler(), Ridge(alpha=lam, fit_intercept=True)])
-    ref = make_pipeline(*steps)
+    ref = make_pipeline(StandardScaler(), Ridge(alpha=lam, fit_intercept=True))
     Y_oh = np.eye(C)[y_tr]
     ref.fit(X_tr, Y_oh)
 
-    # Apply our stored projection (if any) then our ridge weights.
     W = out["W"].cpu().numpy()
     b = out["bias"].cpu().numpy()
-    P = out["projection"]
-    X_proj = X_tr if P is None else X_tr @ P.cpu().numpy().T
-    pred_ours = X_proj @ W + b
+    pred_ours = X_tr @ W + b
     pred_ref = ref.predict(X_tr)
     np.testing.assert_allclose(pred_ours, pred_ref, atol=1e-4)
 
@@ -366,168 +338,88 @@ def test_learner_fit_predict_regression(regression_data):
     assert preds.shape == (len(X_te), 1) or preds.shape == (len(X_te),)
 
 
-def test_max_features_noop_when_D_small(classif_data):
-    """max_features >= D → no projection, result identical to unset."""
-    from open_eeg_bench.ridge_probe import _fit_streaming_ridge
+def test_ridge_probe_consumes_flatten_head_projection(classif_data):
+    """Integration: FlattenHead(max_features=k) caps the features upstream and
+    the ridge probe fits on the k projected dims (lazy buffer materializes
+    during accumulation). Replaces the old probe-level max_features tests."""
+    from open_eeg_bench.head import FlattenHead
+    from open_eeg_bench.ridge_probe import StreamingRidgeProbeLearner
 
     X_tr, y_tr = classif_data["train"]
     X_val, y_val = classif_data["val"]
+    X_te, y_te = classif_data["test"]
     C, D = classif_data["C"], classif_data["D"]
+    k = 64  # < D=200 → projection actually reduces
 
-    out_ref = _fit_streaming_ridge(
-        model=nn.Identity(),
-        train_loader=_loader(X_tr, y_tr),
-        val_loader=_loader(X_val, y_val),
-        n_classes=C,
-        lambdas=[1.0],
-        device="cpu",
-    )
-    out = _fit_streaming_ridge(
-        model=nn.Identity(),
-        train_loader=_loader(X_tr, y_tr),
-        val_loader=_loader(X_val, y_val),
-        n_classes=C,
-        lambdas=[1.0],
-        device="cpu",
-        max_features=D + 10,  # larger than D → projection skipped
-    )
-    assert out["projection"] is None
-    assert out_ref["projection"] is None
-    torch.testing.assert_close(out["W"], out_ref["W"])
-    torch.testing.assert_close(out["bias"], out_ref["bias"])
+    class _Wrap(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.head = nn.Identity()
 
+        def forward(self, x):
+            return self.head(x)
 
-def test_max_features_activates_and_reduces_dim(classif_data):
-    """max_features < D → projection built, n_features matches target."""
-    from open_eeg_bench.ridge_probe import _fit_streaming_ridge
+    model = _Wrap()
+    FlattenHead(max_features=k).apply(model, n_outputs=C, head_module_name="head", seed=0)
 
-    X_tr, y_tr = classif_data["train"]
-    X_val, y_val = classif_data["val"]
-    C, D = classif_data["C"], classif_data["D"]
-    k = 8  # D=20 → project to 8
+    val_set = TensorDataset(torch.from_numpy(X_val), torch.from_numpy(y_val))
+    train_set = TensorDataset(torch.from_numpy(X_tr), torch.from_numpy(y_tr))
+    test_set = TensorDataset(torch.from_numpy(X_te), torch.from_numpy(y_te))
 
-    out = _fit_streaming_ridge(
-        model=nn.Identity(),
-        train_loader=_loader(X_tr, y_tr),
-        val_loader=_loader(X_val, y_val),
+    learner = StreamingRidgeProbeLearner(
+        feature_extractor=model,
         n_classes=C,
-        lambdas=[1.0],
+        batch_size=32,
+        num_workers=0,
         device="cpu",
-        max_features=k,
+        lambdas=[1e-2, 1.0, 1e2],
+        val_set=val_set,
     )
-    P = out["projection"]
-    assert P is not None
-    assert P.shape == (k, D)
-    assert out["n_features"] == k
-    # W is in projected space
-    assert out["W"].shape == (k, C)
-    # Still classifies above chance on separable synthetic data
-    Y_oh = np.eye(C)[y_tr]
-    X64 = torch.from_numpy(X_tr).double()
-    pred = (X64 @ P.T @ out["W"] + out["bias"]).numpy()
-    acc = (pred.argmax(axis=1) == y_tr).mean()
-    assert acc > 1.0 / C
+    learner.fit(train_set, y=None)
+    assert learner._result["n_features"] == k  # ridge saw projected dims
+    assert "projection" not in learner._result  # probe no longer owns projection
+    preds = learner.predict(test_set)
+    assert preds.shape == (len(X_te),)
+    assert (preds == y_te).mean() > 1.0 / C  # above chance
 
 
-def test_max_features_deterministic_across_seeds(classif_data):
-    """Same projection_seed ⇒ same weights; different seed ⇒ different weights."""
+def test_ridge_probe_flatten_head_passthrough_matches_baseline(classif_data):
+    """FlattenHead(max_features >= D) is an identity passthrough: the ridge fit is
+    identical to having no projection at all. Replaces the deleted probe-level
+    test_max_features_noop_when_D_small."""
+    from open_eeg_bench.head import FlattenHead
     from open_eeg_bench.ridge_probe import _fit_streaming_ridge
 
     X_tr, y_tr = classif_data["train"]
     X_val, y_val = classif_data["val"]
-    C = classif_data["C"]
-    kwargs = dict(
-        model=nn.Identity(),
-        train_loader=_loader(X_tr, y_tr),
-        val_loader=_loader(X_val, y_val),
-        n_classes=C,
-        lambdas=[1.0],
-        device="cpu",
-        max_features=8,
-    )
-    out_a = _fit_streaming_ridge(**kwargs, projection_seed=0)
-    out_b = _fit_streaming_ridge(**kwargs, projection_seed=0)
-    out_c = _fit_streaming_ridge(**kwargs, projection_seed=42)
-
-    torch.testing.assert_close(out_a["projection"], out_b["projection"])
-    torch.testing.assert_close(out_a["W"], out_b["W"])
-    assert not torch.allclose(out_a["projection"], out_c["projection"])
-    assert not torch.allclose(out_a["W"], out_c["W"])
-
+    C, D = classif_data["C"], classif_data["D"]
 
-def test_experiment_seed_drives_projection(classif_data):
-    """Experiment.seed flows through Training.build_learner into the projection matrix."""
-    from open_eeg_bench.ridge_probe import StreamingRidgeProbeLearner
-    from open_eeg_bench.training import RidgeProbingTraining
+    class _Wrap(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.head = nn.Identity()
 
-    X_tr, y_tr = classif_data["train"]
-    X_val, y_val = classif_data["val"]
-    C = classif_data["C"]
-    D = classif_data["D"]
+        def forward(self, x):
+            return self.head(x)
 
-    train_set = TensorDataset(torch.from_numpy(X_tr), torch.from_numpy(y_tr))
-    val_set = TensorDataset(torch.from_numpy(X_val), torch.from_numpy(y_val))
-    training_cfg = RidgeProbingTraining(
-        batch_size=32, num_workers=0, device="cpu", max_features=50,
+    model = _Wrap()
+    FlattenHead(max_features=D + 10).apply(  # cap above D → passthrough
+        model, n_outputs=C, head_module_name="head", seed=0
     )
 
-    def fit_with_seed(seed: int) -> StreamingRidgeProbeLearner:
-        learner = training_cfg.build_learner(
-            model=nn.Identity(),
-            callbacks=[],
-            n_classes=C,
-            val_set=val_set,
-            verbose=0,
-            seed=seed,
-        )
-        learner.fit(train_set, y=None)
-        return learner
-
-    learner_a = fit_with_seed(0)
-    learner_b = fit_with_seed(0)
-    learner_c = fit_with_seed(42)
-
-    P_a = learner_a._result["projection"]
-    P_b = learner_b._result["projection"]
-    P_c = learner_c._result["projection"]
-    assert P_a is not None and P_a.shape == (50, D)
-    # Same seed -> same matrix; different seed -> different matrix
-    torch.testing.assert_close(P_a, P_b)
-    assert not torch.allclose(P_a, P_c)
-    # Predictions differ too (matrices differ non-trivially)
-    assert not torch.allclose(learner_a._result["W"], learner_c._result["W"])
-
-
-def test_learner_max_features_predict(classif_data):
-    """Learner applies projection consistently at predict time."""
-    from open_eeg_bench.ridge_probe import StreamingRidgeProbeLearner
-
-    X_tr, y_tr = classif_data["train"]
-    X_val, y_val = classif_data["val"]
-    X_te, y_te = classif_data["test"]
-    C = classif_data["C"]
-
-    train_set = TensorDataset(torch.from_numpy(X_tr), torch.from_numpy(y_tr))
-    val_set = TensorDataset(torch.from_numpy(X_val), torch.from_numpy(y_val))
-    test_set = TensorDataset(torch.from_numpy(X_te), torch.from_numpy(y_te))
-
-    learner = StreamingRidgeProbeLearner(
-        feature_extractor=nn.Identity(),
+    common = dict(
+        train_loader=_loader(X_tr, y_tr),
+        val_loader=_loader(X_val, y_val),
         n_classes=C,
-        batch_size=32,
-        num_workers=0,
+        lambdas=[1.0],
         device="cpu",
-        lambdas=[1e-2, 1.0, 1e2],
-        val_set=val_set,
-        max_features=10,
-        projection_seed=0,
     )
-    learner.fit(train_set, y=None)
-    assert learner._result["projection"] is not None
-    preds = learner.predict(test_set)
-    assert preds.shape == (len(X_te),)
-    # Above chance on synthetic separable data
-    assert (preds == y_te).mean() > 1.0 / C
+    out_proj = _fit_streaming_ridge(model=model, **common)
+    out_ref = _fit_streaming_ridge(model=nn.Identity(), **common)
+
+    assert out_proj["n_features"] == D  # passthrough kept the full dimension
+    torch.testing.assert_close(out_proj["W"], out_ref["W"])
+    torch.testing.assert_close(out_proj["bias"], out_ref["bias"])
 
 
 @pytest.mark.slow