Comparator Model

model

CLASS	DESCRIPTION
ModelComparator	Interface for comparing neural network models.

Classes

ModelComparator(comparator_fn: Comparator | str = 'feature', comparator_kwargs: dict | None = None, matcher_fn: Matcher | str = 'padding', matcher_kwargs: dict | None = None, reduce: Literal['mean', 'median', 'min', 'max'] | None = 'mean')

Interface for comparing neural network models.

This class computes pairwise similarities between all layers of two models, returning an \(N \times M\) comparison matrix where \(N\) and \(M\) are the number of layers in each model.

PARAMETER	DESCRIPTION
comparator_fn	Name of comparison metric to use. Must be registered in ComparatorRegistry. Available: "kernel", "feature", "spectral", "diffusion". TYPE: Comparator | str, optional, by default "feature" DEFAULT: 'feature'
comparator_kwargs	Additional keyword arguments passed to the comparator. TYPE: dict, optional, by default None DEFAULT: None
matcher_fn	Dimension matching strategy: "svd", "projection", "padding", "truncation", "skip" or Matcher. TYPE: Matcher | str, optional, by default "padding" DEFAULT: 'padding'
matcher_kwargs	Additional keyword arguments passed to the matcher. TYPE: dict, optional, by default None DEFAULT: None
reduce	How to reduce the score matrix into a single score. TYPE: Literal["mean", "median", "min", "max"] | None, optional, by default "mean" DEFAULT: 'mean'

METHOD	DESCRIPTION
check_fitted	Check if model has been fitted, raises error if not.
fit	Compare two models by computing pairwise layer similarities.
get_results	Get all fitted parameters as a dictionary.

ATTRIBUTE	DESCRIPTION
score	Returns the scores computed during the fit method. TYPE: Tensor
reduce_score	Returns the reduced scores computed during the fit method. TYPE: Tensor

Source code in spectre/comparator/model.py

def __init__(
    self,
    comparator_fn: Comparator | str = "feature",
    comparator_kwargs: dict | None = None,
    matcher_fn: Matcher | str = "padding",
    matcher_kwargs: dict | None = None,
    reduce: Literal["mean", "median", "min", "max"] | None = "mean",
):
    self.comparator_fn = initialize_comparator_fn(
        comparator_fn=comparator_fn,
        comparator_kwargs=comparator_kwargs or {},
    )

    self.matcher_fn = initialize_matcher_fn(
        matcher_fn=matcher_fn,
        matcher_kwargs=matcher_kwargs or {},
    )

    allowed_reduce = ["mean", "median", "min", "max", None]
    if reduce not in allowed_reduce:
        raise ValueError(
            f"Invalid reduce '{reduce}'. Choose from: {allowed_reduce}"
        )
    self.reduce = reduce

    # Internal state
    self._score = None
    self._is_fitted = False

Attributes

score: torch.Tensor property

Returns the scores computed during the fit method.

reduce_score: torch.Tensor property

Returns the reduced scores computed during the fit method.

Functions

check_fitted() -> None

Check if model has been fitted, raises error if not.

Source code in spectre/comparator/model.py

def check_fitted(self) -> None:
    """Check if model has been fitted, raises error if not."""
    if not self._is_fitted:
        raise ValueError("Model has not been fitted yet.")

fit(model_1: torch.nn.Module | Parametric, model_2: torch.nn.Module | Parametric, X: torch.Tensor, Y: torch.Tensor | None = None, layers_1: list[str] | None = None, layers_2: list[str] | None = None) -> ModelComparator

Compare two models by computing pairwise layer similarities.

This method computes an \(N imes M\) comparison matrix where \(N\) is the number of layers in model_1 and \(M\) is the number of layers in model_2. Each entry (i, j) contains the similarity score between layer i of model_1 and layer j of model_2.

Optionally, if Y is provided, the method computes the similarity between the output of model_1 on input X and model_2 on the input Y.

PARAMETER	DESCRIPTION
model_1	First model (n_layers length). TYPE: Module | Parametric
model_2	Second model (m_layers length). TYPE: Module | Parametric
X	Input data to pass through both models if Y is not provided. TYPE: Tensor
Y	Input data to pass through model_2. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
layers_1	Layer names to extract from model_1. If None, use all leaf modules. TYPE: list[str] | None, optional, by default None DEFAULT: None
layers_2	Layer names to extract from model_2. If None, use all leaf modules. TYPE: list[str] | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
ModelComparator	Self for chaining.

Examples:

>>> comparator = ModelComparator("kernel", matcher="svd", reduce=None)
>>> comparator.fit(model1, model2, X)
>>> comparator.score.shape
torch.Size([10, 12])  # model1 has 10 layers, model2 has 12 layers
>>> result = comparator.get_results()
>>> result["features_1"].keys()
dict_keys(['conv1', 'relu1', 'conv2', ...])
>>> result["features_2"].keys()
dict_keys(['layer1', 'layer2', ...])

With reduction:

>>> comparator = ModelComparator("kernel", reduce="mean")
>>> comparator.fit(model1, model2, X)
>>> comparator.score  # Reduced score
tensor(0.823)

Source code in spectre/comparator/model.py

def fit(
    self,
    model_1: torch.nn.Module | Parametric,
    model_2: torch.nn.Module | Parametric,
    X: torch.Tensor,
    Y: torch.Tensor | None = None,
    layers_1: list[str] | None = None,
    layers_2: list[str] | None = None,
) -> "ModelComparator":
    """
    Compare two models by computing pairwise layer similarities.

    This method computes an $N \times M$ comparison matrix where $N$ is the
    number of layers in `model_1` and $M$ is the number of layers in `model_2`.
    Each entry `(i, j)` contains the similarity score between layer `i` of `model_1`
    and layer `j` of `model_2`.

    Optionally, if `Y` is provided, the method computes the similarity between
    the output of `model_1` on input `X` and `model_2` on the input `Y`.

    Parameters
    ----------
    model_1 : torch.nn.Module | Parametric
        First model (`n_layers` length).

    model_2 : torch.nn.Module | Parametric
        Second model (`m_layers` length).

    X : torch.Tensor
        Input data to pass through both models if `Y` is not provided.

    Y : torch.Tensor | None, optional, by default None
        Input data to pass through `model_2`.

    layers_1 : list[str] | None, optional, by default None
        Layer names to extract from `model_1`. If None, use all leaf modules.

    layers_2 : list[str] | None, optional, by default None
        Layer names to extract from `model_2`. If None, use all leaf modules.

    Returns
    -------
    ModelComparator
        Self for chaining.

    Examples
    --------
    >>> comparator = ModelComparator("kernel", matcher="svd", reduce=None)
    >>> comparator.fit(model1, model2, X)
    >>> comparator.score.shape
    torch.Size([10, 12])  # model1 has 10 layers, model2 has 12 layers
    >>> result = comparator.get_results()
    >>> result["features_1"].keys()
    dict_keys(['conv1', 'relu1', 'conv2', ...])
    >>> result["features_2"].keys()
    dict_keys(['layer1', 'layer2', ...])

    With reduction:

    >>> comparator = ModelComparator("kernel", reduce="mean")
    >>> comparator.fit(model1, model2, X)
    >>> comparator.score  # Reduced score
    tensor(0.823)
    """
    Y = X if Y is None else Y

    self._features_1 = extract_features(model=model_1, X=X, layers=layers_1)
    self._features_2 = extract_features(model=model_2, X=Y, layers=layers_2)

    layers_1_list = list(self._features_1.keys())
    layers_2_list = list(self._features_2.keys())

    self._score = torch.full(
        (len(layers_1_list), len(layers_2_list)),
        float("nan"),
        dtype=X.dtype,
        device=X.device,
    )

    for i, l_1 in enumerate(layers_1_list):
        for j, l_2 in enumerate(layers_2_list):
            X1 = self._features_1[l_1]
            X2 = self._features_2[l_2]

            try:
                X1_matched, X2_matched = self.matcher_fn(X1, X2)
            except ValueError:
                # Skip pairs that can't be matched (if matcher="skip")
                continue

            self._score[i, j] = self.comparator_fn(X1_matched, X2_matched)

    if torch.isnan(self._score).all():
        warnings.warn(
            "All layer pairs failed to match. No similarities could be computed. "
            f"Use a different matcher (current: '{self.matcher_fn.name}').",
            UserWarning,
        )

    if self.reduce is not None:
        # Only reduce over non-NaN values
        allowed_scores = self._score[~torch.isnan(self._score)]
        if len(allowed_scores) > 0:
            if self.reduce == "mean":
                self._reduce_score = torch.mean(allowed_scores)
            elif self.reduce == "median":
                self._reduce_score = torch.median(allowed_scores)
            elif self.reduce == "min":
                self._reduce_score = torch.min(allowed_scores)
            elif self.reduce == "max":
                self._reduce_score = torch.max(allowed_scores)
        else:
            warnings.warn("No valid scores to reduce.", UserWarning)
            self._reduce_score = None
    else:
        self._reduce_score = None

    self._is_fitted = True

    return self

get_results() -> dict[str, torch.Tensor | dict[str, torch.Tensor] | str | None]

Get all fitted parameters as a dictionary.

RETURNS

DESCRIPTION

dict[str, Tensor | dict[str, Tensor] | str | None]

Result dictionary with keys: "score": torch.Tensor, of shape (n_layers, m_layers). Pairwise similarity scores between all layer pairs. "reduce": float | None. Reduced score if reduce is not None "features_1": dict[str, torch.Tensor]. Features extracted from model_1 "features_2": dict[str, torch.Tensor]. Features extracted from model_2 "matcher": str. Name of the matcher used "comparator_name": str. Name of the comparator used

Source code in spectre/comparator/model.py

def get_results(
    self,
) -> dict[str, torch.Tensor | dict[str, torch.Tensor] | str | None]:
    """
    Get all fitted parameters as a dictionary.

    Returns
    -------
    dict[str, torch.Tensor | dict[str, torch.Tensor] | str | None]
        Result dictionary with keys:

        - "score": torch.Tensor, of shape (`n_layers`, `m_layers`). Pairwise
          similarity scores between all layer pairs.
        - "reduce": float | None. Reduced score if `reduce` is not None
        - "features_1": dict[str, torch.Tensor]. Features extracted from `model_1`
        - "features_2": dict[str, torch.Tensor]. Features extracted from `model_2`
        - "matcher": str. Name of the matcher used
        - "comparator_name": str. Name of the comparator used
    """
    self.check_fitted()

    return {
        "score": self._score,
        "reduce": self._reduce_score,
        "features_1": self._features_1,
        "features_2": self._features_2,
        "matcher": self.matcher_fn.name,
        "comparator": self.comparator_fn.name,
    }

Functions