[WIP] Lazy projection

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

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[

open_eeg_bench/head_utils.py

@@ -0,0 +1,115 @@
+"""Lazy random-projection head utilities.
+
+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 numpy as np
+import torch
+import torch.nn as nn
+from torch.nn.modules.lazy import LazyModuleMixin
+from torch.nn.parameter import UninitializedBuffer
+
+
+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 that caps the feature dimension lazily.
+
+    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: 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
+            :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_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.copy_(P)
+
+    def forward(self, input):
+        return input @ self.projection_matrix.T