Skip to content

mlvbench.probes

Probes.

Probes are small models that readout information from intermediate features of a pretrained vision model.

LinearProbe

Bases: Probe

Linear probe.

A linear projection of the input features to a scalar, optionally with a learnable bias:

\[ y = W x + b \]

name property

Name of the probe type ("linear").

__init__(layer, input_dim, bias='scalar', num_patches=None, metadata=None)

Initialize the probe.

Parameters:

Name Type Description Default
layer str

Name of the model layer the probe reads from.

 required
input_dim int

Number of dimensions of the input features.

 required
bias Literal['none', 'scalar', 'spatial'] | Tensor

The kind of learnable bias added to the output:

  • "none": no bias.
  • "scalar": a single learnable scalar shared across all patches.
  • "spatial": a learnable per-patch bias with shape (num_patches, 1), initialized to zero. Requires num_patches.
  • a torch.Tensor with shape (N, 1): a per-patch bias initialized from the tensor; equivalent to "spatial" with num_patches = N.
 'scalar'
num_patches int | None

Number of patch tokens of the input features. Required when using a "spatial" bias and ignored otherwise.

 None
metadata dict[str, Any] | None

Free-form metadata to attach to the probe.

 None

forward(features)

Forward pass.

Parameters:

Name Type Description Default
features Features

The features extracted from the model. The probe uses the patch tokens of self.layer, expecting shape (B, N, C) and dtype float32.

 required

Returns:

Type Description
Tensor

The output with shape (B, N, 1) and dtype float32.

select_features(features)

Select the part of the feature representation used by the probe.

Defaults to the patch tokens of self.layer with shape (B, N, C). Override to read from a different part of the representation (e.g. the cls token or a different token format).

Parameters:

Name Type Description Default
features Features

The features extracted from the backbone.

 required

Returns:

Type Description
Tensor

The selected feature tensor.

serialize()

Serialize the probe.

See Probe.serialize.

Probe

Bases: Module, ABC

Abstract base class for probes.

Attributes:

Name Type Description
name str

Name of the probe type.
layer str

Name of the backbone layer the probe reads from.
metadata dict[str, Any]

Free-form dict of training/experiment information attached to the probe. Persisted in full when the probe is saved.

name abstractmethod property

Name of the probe type.

__init__(layer, input_dim, *, metadata=None)

Initialize the probe.

Parameters:

Name Type Description Default
layer str

Name of the model layer the probe reads from.

 required
input_dim int

Number of dimensions of the input features.

 required
metadata dict[str, Any] | None

Free-form dict of training/experiment information.

 None

forward(features) abstractmethod

Apply the probe to the given features.

The probe selects the part of the representation it reads from via select_features.

Parameters:

Name Type Description Default
features Features

The features extracted from the backbone.

 required

Returns:

Type Description
Tensor

The prediction from the probe as torch.Tensor. Shape and dtype very per probe.

select_features(features)

Select the part of the feature representation used by the probe.

Defaults to the patch tokens of self.layer with shape (B, N, C). Override to read from a different part of the representation (e.g. the cls token or a different token format).

Parameters:

Name Type Description Default
features Features

The features extracted from the backbone.

 required

Returns:

Type Description
Tensor

The selected feature tensor.

serialize() abstractmethod

Return all information needed to reconstruct the probe.

This returns a self-describing dictionary that includes both the torch state_dict and all config options. The probe can be rebuilt from this representation via build_probe.

Returns:

Type Description
dict[str, Any]

A dictionary describing the probe, with detached CPU tensors for weights.

build_probe(name, layer, input_dim, metadata=None, state_dict=None, probe_types=None, **kwargs)

Build the probe with the given name and arguments.

Parameters:

Name Type Description Default
name str

Name of the probe type (e.g. "linear").

 required
layer str

Name of the model layer the probe reads from.

 required
input_dim int

Number of dimensions of the input features.

 required
metadata dict[str, Any] | None

Free-form metadata to attach to the probe.

 None
state_dict dict[str, Any] | None

Optional state dict for restoring pretrained probes.

 None
probe_types dict[str, type[Probe]] | None

Optional mapping of custom probe type names to Probe subclasses, merged with the built-in types. Each class must be constructible as Cls(layer, input_dim, metadata=metadata, **kwargs).

 None
**kwargs Any

Additional keyword arguments forwarded to the probe constructor (e.g. bias and num_patches).

 {}

Raises:

Type Description
ValueError

If name is not a known built-in or custom probe type.