← Back to API reference

neuralbench.modules.DownstreamWrapper

pydantic model neuralbench.modules.DownstreamWrapper

Configuration for wrapping a (pretrained) model for downstream fine-tuning or linear probing.

This class provides a declarative way to configure how a pretrained model should be adapted for downstream tasks, including optional on-the-fly preprocessing, layer freezing, output aggregation, and adding a trainable probe on top of the model.

Parameters:

• on_the_fly_preprocessor (OnTheFlyPreprocessor | None, optional) – On-the-fly preprocessing applied to the input before the model forward pass. Typically model-specific (e.g. QuantileAbsScaler for BIOT). Default is None.

• channel_adapter_config (ChannelMerger | ChannelProjection | None, optional) – Configuration for a channel adapter that projects from arbitrary input channels to a fixed number of target channels. Supply a ChannelMerger for position-based spatial attention, or a ChannelProjection for a simple Conv1d(kernel_size=1) linear mixing. Default is None.

• model_output_key (str | int | None, optional) – Key or index to extract from model output dictionary. If None, assumes the model returns a tensor directly. Default is None.

• layers_to_freeze (list[str] | None, optional) – List of layer name patterns to freeze (set requires_grad=False). Cannot be used together with layers_to_unfreeze. Default is None.

• layers_to_unfreeze (list[str] | tp.Literal["last"] | None, optional) – List of layer name patterns to unfreeze (set requires_grad=True), while freezing all others. Cannot be used together with layers_to_freeze. If “last”, unfreezes the last layer (nn.Module) of the model. Default is None.

• strict_matching (bool, optional) – If True, when freezing/unfreezing layers, only the first part of the layer name (before the first dot) must match exactly. If False, any part of the layer name can match the patterns. Default is True.

• aggregation ({"flatten", "mean", "first"} or int, optional) – Parameter-free reduction of the model output before the probe. "flatten" flattens all dimensions except batch; "mean" averages over the temporal/sequence dimension (dim=1); "first" selects only the first timestep/token; an int splits into n groups, averages each group, then concatenates; None performs no aggregation (required by probe_config="attention"). When probe_layer is set, the captured activation is canonicalised to batch-first before aggregation (see probe_batch_dim), so these semantics are identical for intermediate and final outputs.

• probe_config (Mlp | "linear" | "attention" | None, optional) – Configuration for the probe head added on top. None uses identity (no additional layer), e.g. if the model already has a linear layer of the right output size. "linear" adds a single linear layer. "attention" adds an AttentivePool read-out plus a linear layer; it pools the raw tokens itself, so requires aggregation=None. An Mlp instance adds a multi-layer perceptron with specified configuration.

• probe_layer (str | None, optional) – Dotted submodule name (from model.named_modules()) where a forward hook taps activations for probing. None (default) probes the final model output. Requires model_output_key=None (intermediate captures are tensors, not dicts).

• probe_batch_dim (int | "auto", optional) – Axis of the probed activation that indexes the batch. "auto" (default) detects it by running the dummy forward at two batch sizes and finding the axis that scales with the batch. Set explicitly (e.g. 1 for sequence-first (T, B, D) transformer outputs) to skip detection or resolve an ambiguous layout. Only used when probe_layer is set.

• lora_config (LoraConfig | None, optional) – If set, PEFT LoRA adapters are injected into the wrapped model, leaving only the adapters and the probe head trainable. Requires the peft package. Default is None.

• lora_target_modules (list[str] | None, optional) – nn.Linear leaf-module names the adapters target (e.g. ["to_q", "to_k", "to_v", "to_out"]), usually set per foundation model in its YAML. Overridden by lora_config.target_modules, ignored when lora_config is None. Default is None.

Fields:

• aggregation (Literal['flatten', 'mean', 'first'] | int | None)

• channel_adapter_config (neuraltrain.models.common.ChannelMerger | neuralbench.modules.ChannelProjection | None)

• layers_to_freeze (list[str] | None)

• layers_to_unfreeze (list[str] | Literal['last'] | None)

• lora_config (neuralbench.modules.LoraConfig | None)

• lora_target_modules (list[str] | None)

• model_output_key (str | int | None)

• on_the_fly_preprocessor (neuraltrain.models.preprocessor.OnTheFlyPreprocessor | None)

• probe_batch_dim (int | Literal['auto'])

• probe_config (neuraltrain.models.common.Mlp | Literal['linear', 'attention'] | None)

• probe_layer (str | None)

• strict_matching (bool)

field on_the_fly_preprocessor: OnTheFlyPreprocessor | None = None

field channel_adapter_config: ChannelMerger | ChannelProjection | None = None

field model_output_key: str | int | None = None

field layers_to_freeze: list[str] | None = None

field layers_to_unfreeze: list[str] | Literal['last'] | None = None

field strict_matching: bool = True

field aggregation: Literal['flatten', 'mean', 'first'] | int | None = 'flatten'

field probe_config: Mlp | Literal['linear', 'attention'] | None = 'linear'

field probe_layer: str | None = None

field probe_batch_dim: int | Literal['auto'] = 'auto'

field lora_config: LoraConfig | None = None

field lora_target_modules: list[str] | None = None

property n_adapter_target_channels: int | None

Target channel count of the adapter, or None if no adapter is configured.

build(model: Module, dummy_batch: dict[str, Tensor | None], n_outputs: int, input_channel_names: list[str] | None = None) → DownstreamWrapperModel

← Back to API reference