Fix backbone layers (#25)

* Use dev braindecode (will remove later)

* Remove model_kwargs from backbone configs (kwargs are coming from hf_hub

* Update sJEPA hub repo

* Raise errors, not warnings when the backbones are not correctly configured

* Add Labram's temporal_embedding to training_required_modules (does not work yet because temporal_embedding is a parameter, not a module)

* Update uv loc for latest braindecode dev version

* Use new Interpolated* backbones

* update braindecode

* Add training_required_parameters field on _BackboneBase

Distinct from training_required_modules: this one accepts top-level
nn.Parameter names (not Modules), e.g. Labram's temporal_embedding.
Extend _check_layers_and_parameters_exist to validate against
named_parameters(), and broaden the checkpoint loading allowed_names
so missing-param entries can be covered by either field.

* Add supports_training_required_parameters ClassVar and wire PEFT target_parameters

Each finetuning method declares (via a ClassVar, not a model field, so
it is immutable from the constructor under extra='forbid') whether it
can keep top-level nn.Parameters trainable. LoRA/AdaLoRA/DoRA forward
training_required_parameters to PEFT's target_parameters argument;
Frozen/TwoStages add the parameter names to the unfreeze list. IA3 and
OFT opt out (no target_parameters in their PEFT configs).

* Validate finetuning compatibility with training_required_parameters

The Experiment validator now reads
finetuning.supports_training_required_parameters and rejects mismatched
combinations. Also extend the ridge check to reject non-empty
training_required_parameters (ridge cannot train extra parameters).

* Migrate Labram default config to InterpolatedLaBraM and use training_required_parameters

Completes the migration started in 3b4ed0e (which migrated BENDR and
SignalJEPA). InterpolatedLaBraM accepts arbitrary montages by projecting
onto the canonical 128-channel layout. temporal_embedding (still an
nn.Parameter) now declared via training_required_parameters.

* Add tests for training_required_parameters validation

Three tests mirroring test_ridge_rejects_training_required_modules:
ridge rejects training_required_parameters, IA3 rejects it, LoRA
accepts it (LoRA forwards to PEFT target_parameters).

* Document training_required_parameters in CHANGELOG

* fix warning message

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: PierreGtch

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Allow custom normalization methods by subclassing `Normalization`, now based on `exca.helpers.DiscriminatedModel`. Builtin subclasses pin their pre-existing `kind` value to preserve cached experiment UIDs ([#35](https://github.com/braindecode/OpenEEGBench/pull/35)).
+- Add `training_required_parameters` field on `_BackboneBase` for top-level `nn.Parameter` names that must remain trainable (e.g. Labram's `temporal_embedding`). Distinct from `training_required_modules`, which only accepts `nn.Module` names. Finetuning methods declare compatibility via the `supports_training_required_parameters` class variable; `IA3` and `OFT` opt out. `LoRA`/`AdaLoRA`/`DoRA` forward the list to PEFT's `target_parameters`; `Frozen`/`TwoStages` extend their unfreeze list ([#25](https://github.com/braindecode/OpenEEGBench/pull/25)).
 
 ## [0.4.0] - 2026-05-07
 

open_eeg_bench/backbone.py

@@ -18,7 +18,6 @@
 import logging
 from typing import Any, Literal
 from importlib import import_module
-import warnings
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 
@@ -58,12 +57,24 @@ class _BackboneBase(BaseModel):
             "since training extra layers is not comparable to training only a classification head."
         ),
     )
+    training_required_parameters: list[str] | None = Field(
+        default=None,
+        description=(
+            "Top-level ``nn.Parameter`` names (NOT modules) that must remain trainable, "
+            "typically because they cannot be loaded from the pretrained checkpoint "
+            "(e.g. shape mismatch with the new ``n_chans``/``n_times``). "
+            "Only finetuning methods whose "
+            "``supports_training_required_parameters`` is True accept a non-empty list. "
+            "Like ``training_required_modules``, models using this field are categorized "
+            "separately in evaluations."
+        ),
+    )
     normalization: Normalization = Field(
         default=NoNormalization(),
         description="Post-window normalization applied to each data window.",
     )
 
-    def get_training_required_modules(self):
+    def get_training_required_modules(self) -> list[str]:
         out = [self.head_module_name]
         if self.training_required_modules:
             log.warning(
@@ -75,6 +86,16 @@ def get_training_required_modules(self):
             out += self.training_required_modules
         return out
 
+    def get_training_required_parameters(self) -> list[str]:
+        if self.training_required_parameters:
+            log.warning(
+                "%s requires training extra parameters %s beyond the head. "
+                "This model will be categorized separately in evaluations.",
+                type(self).__name__,
+                self.training_required_parameters,
+            )
+        return list(self.training_required_parameters or [])
+
     def _build(
         self,
         n_chans: int,
@@ -92,6 +113,7 @@ def _check_layers_and_parameters_exist(self, model):
         This test can only be done once the model has been instantiated.
         """
         module_names = {name for name, _ in model.named_modules()}
+        param_names = {name for name, _ in model.named_parameters()}
         module_fields = {"head_module_name": [self.head_module_name]}
         if (
             self.peft_target_modules is not None
@@ -102,16 +124,26 @@ def _check_layers_and_parameters_exist(self, model):
             module_fields["peft_ff_modules"] = self.peft_ff_modules
         if self.training_required_modules is not None:
             module_fields["training_required_modules"] = self.training_required_modules
-        for field, names in module_fields.items():
-            for name in names:
-                # Allow short names that match a suffix (e.g. "qkv" matches "encoder.layer.0.qkv")
-                if name not in module_names and not any(
-                    n == name or n.endswith("." + name) for n in module_names
-                ):
-                    raise ValueError(
-                        f"{type(self).__name__}.{field} references module '{name}' "
-                        f"which does not exist in the model."
-                    )
+        param_fields = {}
+        if self.training_required_parameters is not None:
+            param_fields["training_required_parameters"] = (
+                self.training_required_parameters
+            )
+
+        def _check(fields, valid_names, kind):
+            for field, names in fields.items():
+                for name in names:
+                    # Allow short names matching a suffix (e.g. "qkv" matches "encoder.layer.0.qkv")
+                    if name not in valid_names and not any(
+                        n == name or n.endswith("." + name) for n in valid_names
+                    ):
+                        raise ValueError(
+                            f"{type(self).__name__}.{field} references {kind} '{name}' "
+                            f"which does not exist in the model."
+                        )
+
+        _check(module_fields, module_names, "module")
+        _check(param_fields, param_names, "parameter")
 
     def build(
         self,
@@ -277,39 +309,60 @@ def load_pretrained(self, model) -> None:
             len(skipped),
         )
 
-        # Sanity check on missing keys (model keys not loaded from checkpoint).
-        param_names = {name for name, _ in model.named_parameters()}
-        allowed_prefixes = [self.head_module_name]
-        if self.training_required_modules:
-            allowed_prefixes.extend(self.training_required_modules)
+        allowed_names = (
+            [self.head_module_name]
+            + (self.training_required_modules or [])
+            + (self.training_required_parameters or [])
+        )
+
+        def covered(k: str) -> bool:
+            return any(f".{m}." in f".{k}." for m in allowed_names)
 
-        def _is_allowed(key: str) -> bool:
-            return any(key == p or key.startswith(p + ".") for p in allowed_prefixes)
+        def describe(k: str) -> str:
+            if k in state_dict:
+                ckpt_shape = tuple(state_dict[k].shape)
+                model_shape = tuple(model_state[k].shape)
+                return f"{k}  [shape mismatch: ckpt {ckpt_shape} vs model {model_shape}]"
+            return f"{k}  [absent from checkpoint]"
 
-        unexpected_params = [
-            k for k in missing if k in param_names and not _is_allowed(k)
+        param_names = {n for n, _ in model.named_parameters()}
+        missing_params = [k for k in missing if k in param_names and not covered(k)]
+        missing_buffers = [
+            k for k in missing if k not in param_names and not covered(k)
         ]
-        missing_buffers = [k for k in missing if k not in param_names]
-
-        if unexpected_params:
-            # TODO: 
-            # - FIX the models. 
-            # - Transform this warning into an error
-            warnings.warn(
-                f"Pretrained checkpoint is missing weights for backbone "
-                f"parameters outside of head_module_name and "
-                f"training_required_modules: {unexpected_params}. "
-                f"These parameters keep random initialization, making results "
-                f"seed-dependent. Either provide a checkpoint that covers them "
-                f"or add the relevant modules to training_required_modules.",
-                stacklevel=2,
+
+        if missing_params:
+            raise ValueError(
+                f"Pretrained checkpoint for {self.model_cls} is missing "
+                f"{len(missing_params)} learnable parameter(s) that are neither "
+                f"under head_module_name='{self.head_module_name}' nor under any "
+                f"name in training_required_modules={self.training_required_modules} "
+                f"or training_required_parameters={self.training_required_parameters}:\n"
+                + "\n".join(f"  - {describe(k)}" for k in missing_params)
+                + "\nThese parameters would be silently initialized from scratch, "
+                "which is almost certainly a config error. Either:\n"
+                "  (a) the checkpoint is incompatible with the declared architecture "
+                "(check model_kwargs), or\n"
+                "  (b) these names should be declared in `training_required_modules` "
+                "(for nn.Modules) or `training_required_parameters` (for top-level "
+                "nn.Parameters) so they are explicitly trained from scratch (note: "
+                "this will categorize the model separately in evaluations)."
             )
+
         if missing_buffers:
-            warnings.warn(
-                f"Pretrained checkpoint does not cover the following "
-                f"buffers: {missing_buffers}. They keep their default "
-                f"initialization (typically deterministic).",
-                stacklevel=2,
+            log.warning(
+                "Pretrained checkpoint for %s is missing %d buffer(s) that are "
+                "not covered by the allowlist "
+                "(head_module_name='%s', training_required_modules=%s, "
+                "training_required_parameters=%s):\n%s\n"
+                "Buffers are not trained, so missing values may be computed at "
+                "init — but verify this is intentional.",
+                self.model_cls,
+                len(missing_buffers),
+                self.head_module_name,
+                self.training_required_modules,
+                self.training_required_parameters,
+                "\n".join(f"  - {describe(k)}" for k in missing_buffers),
             )
 
     @staticmethod

open_eeg_bench/default_configs/backbones.py

@@ -12,17 +12,6 @@
 def biot(**overrides) -> PretrainedBackbone:
     defaults = dict(
         model_cls="braindecode.models.BIOT",
-        model_kwargs=dict(
-            embed_dim=256,
-            num_heads=8,
-            num_layers=4,
-            drop_prob=0.5,
-            att_drop_prob=0.2,
-            att_layer_drop_prob=0.2,
-            hop_length=100,
-            max_seq_len=1024,
-            return_feature=False,
-        ),
         # w1, w2 are the two Linear layers inside FeedForward blocks
         peft_ff_modules=["w1", "w2"],
         normalization=PercentileScale(q=95.0),
@@ -35,26 +24,12 @@ def biot(**overrides) -> PretrainedBackbone:
 
 def labram(**overrides) -> PretrainedBackbone:
     defaults = dict(
-        model_cls="braindecode.models.Labram",
-        model_kwargs=dict(
-            patch_size=200,
-            embed_dim=200,
-            num_layers=12,
-            num_heads=10,
-            mlp_ratio=4.0,
-            qkv_bias=False,
-            qk_scale=None,
-            drop_prob=0.1,
-            attn_drop_prob=0.0,
-            drop_path_prob=0.1,
-            use_abs_pos_emb=True,
-            use_mean_pooling=True,
-            init_scale=0.001,
-            init_values=0.1,
-            neural_tokenizer=True,
-        ),
+        model_cls="braindecode.models.InterpolatedLaBraM",
         # mlp.0 and mlp.2 are the two Linear layers inside the MLP block
         peft_ff_modules=["mlp.0", "mlp.2"],
+        # temporal_embedding is an nn.Parameter trained for n_times=3000 only
+        # (shape mismatch with other n_times → must be retrained from scratch)
+        training_required_parameters=["temporal_embedding"],
         normalization=DivideByConstant(factor=100.0),
         hub_repo="braindecode/labram-pretrained",
     )
@@ -64,26 +39,9 @@ def labram(**overrides) -> PretrainedBackbone:
 
 def bendr(**overrides) -> PretrainedBackbone:
     defaults = dict(
-        model_cls="braindecode.models.BENDR",
-        model_kwargs=dict(
-            encoder_h=512,
-            contextualizer_hidden=3076,
-            transformer_layers=8,
-            transformer_heads=8,
-            position_encoder_length=25,
-            enc_width=(3, 2, 2, 2, 2, 2),
-            enc_downsample=(3, 2, 2, 2, 2, 2),
-            drop_prob=0.1,
-            layer_drop=0.0,
-            projection_head=False,
-            start_token=-5,
-            final_layer=True,
-            n_chans_pretrained=20,
-            encoder_only=True,
-        ),
+        model_cls="braindecode.models.InterpolatedBENDR",
         # linear1, linear2 are the FFN layers in TransformerEncoderLayer
         peft_ff_modules=["linear1", "linear2"],
-        training_required_modules=["channel_projection"],
         normalization=MinMaxScale(),
         hub_repo="braindecode/braindecode-bendr",
     )
@@ -94,20 +52,6 @@ def bendr(**overrides) -> PretrainedBackbone:
 def cbramod(**overrides) -> PretrainedBackbone:
     defaults = dict(
         model_cls="braindecode.models.CBraMod",
-        model_kwargs=dict(
-            patch_size=200,
-            dim_feedforward=800,
-            n_layer=12,
-            nhead=8,
-            emb_dim=200,
-            drop_prob=0.1,
-            channels_kernel_stride_padding_norm=[
-                [25, 49, 25, 24, [5, 25]],
-                [25, 3, 1, 1, [5, 25]],
-                [25, 3, 1, 1, [5, 25]],
-            ],
-            return_encoder_output=False,
-        ),
         # linear1, linear2 are the FFN layers in CrissCrossTransformerEncoderLayer
         peft_ff_modules=["linear1", "linear2"],
         normalization=DivideByConstant(factor=100.0),
@@ -119,29 +63,10 @@ def cbramod(**overrides) -> PretrainedBackbone:
 
 def signal_jepa(**overrides) -> PretrainedBackbone:
     defaults = dict(
-        model_cls="braindecode.models.SignalJEPA",
-        model_kwargs=dict(
-            feature_encoder__conv_layers_spec=[
-                (8, 32, 8),
-                (16, 2, 2),
-                (32, 2, 2),
-                (64, 2, 2),
-                (64, 2, 2),
-            ],
-            feature_encoder__mode="default",
-            feature_encoder__conv_bias=False,
-            pos_encoder__spat_dim=30,
-            pos_encoder__time_dim=34,
-            pos_encoder__sfreq_features=1.0,
-            transformer__d_model=64,
-            transformer__num_encoder_layers=8,
-            transformer__num_decoder_layers=4,
-            transformer__nhead=8,
-            drop_prob=0.0,
-        ),
+        model_cls="braindecode.models.InterpolatedSignalJEPA",
         # linear1, linear2 are the FFN layers in TransformerEncoderLayer
         peft_ff_modules=["linear1", "linear2"],
-        checkpoint_url="https://huggingface.co/braindecode/SignalJEPA/resolve/main/signal-jepa_16s-60_adeuwv4s.pth",
+        hub_repo="braindecode/signal-jepa",
     )
     defaults.update(overrides)
     return PretrainedBackbone(**defaults)
@@ -150,17 +75,6 @@ def signal_jepa(**overrides) -> PretrainedBackbone:
 def reve(**overrides) -> PretrainedBackbone:
     defaults = dict(
         model_cls="braindecode.models.REVE",
-        model_kwargs=dict(
-            embed_dim=512,
-            depth=22,
-            heads=8,
-            head_dim=64,
-            mlp_dim_ratio=2.66,
-            use_geglu=True,
-            patch_size=200,
-            patch_overlap=20,
-            attention_pooling=False,
-        ),
         # net.1, net.3 are the two Linear layers inside FeedForward.net
         peft_ff_modules=["net.1", "net.3"],
         normalization=WindowZScore(clip_sigma=15.0),

open_eeg_bench/experiment.py

@@ -79,13 +79,31 @@ def _check_consistency(self):
         if is_ridge and not isinstance(self.finetuning, Frozen):
             raise ValueError("Ridge probing requires Frozen finetuning.")
 
-        if is_ridge and self.backbone.training_required_modules:
+        if is_ridge and (
+            self.backbone.training_required_modules
+            or self.backbone.training_required_parameters
+        ):
             raise ValueError(
                 f"Ridge probing incompatible with backbone.training_required_modules="
-                f"{self.backbone.training_required_modules}. "
+                f"{self.backbone.training_required_modules} or "
+                f"backbone.training_required_parameters="
+                f"{self.backbone.training_required_parameters}. "
                 f"These backbones need SGD training."
             )
 
+        if (
+            self.backbone.training_required_parameters
+            and not self.finetuning.supports_training_required_parameters
+        ):
+            raise ValueError(
+                f"{type(self.finetuning).__name__} is incompatible with "
+                f"backbone.training_required_parameters="
+                f"{self.backbone.training_required_parameters}. "
+                f"This finetuning method can only train nn.Modules. "
+                f"Use a method whose supports_training_required_parameters is True "
+                f"(e.g. LoRA, AdaLoRA, DoRA, FullFinetune, Frozen, TwoStages)."
+            )
+
         if (
             not is_ridge
             and isinstance(self.finetuning, Frozen)