Loss Distillation

distillation

CLASS	DESCRIPTION
DistillationLoss	Distillation loss that adapts to teacher-student pair.
FeatureMatchingLoss	Loss for matching intermediate feature representations.
SpectralPreservationLoss	Loss for preserving spectral properties in distillation.

FUNCTION	DESCRIPTION
distillation_loss	Compute combined distillation loss and task loss (MSE by default).
feature_matching_loss	Compute feature matching loss between teacher and student features.
spectral_preservation_loss	Loss for preserving spectral properties in distillation.

Classes

DistillationLoss(alpha: float = 0.7, temperature: float = 4.0, shape_matching: Literal['projection', 'padding', 'truncation'] = 'projection', reduce: Literal['sum', 'mean'] = 'mean', sign: float = 1.0, context_prefix: str | None = None)

Bases: Loss

Distillation loss that adapts to teacher-student pair.

Combines distillation loss (student vs teacher) with task loss (student vs ground truth) with using a weighted combination.

Loss is given as \(\alpha L_D + (1 - \alpha) L_T\), where \(\alpha\) balances the two components.

Distillation loss is given as \(L_D = t^2 \mathrm{KL}(\text{softmax}(S / t)~\|~\text{softmax}(T / t))\), where \(t\) is the temperature parameter.

Task loss is given as mean squared error \(L_T = \text{MSE}(S, Y)\), where \(S\) is the student output and \(Y\) is the ground truth.

PARAMETER	DESCRIPTION
alpha	Weight for distillation loss vs task loss. alpha=1.0: Pure distillation (no task loss) alpha=0.0: Pure task loss (no distillation) 0 < alpha < 1: Weighted combination TYPE: float, optional, by default 0.7 DEFAULT: 0.7
temperature	Temperature for soft target scaling. TYPE: float, optional, by default 4.0 DEFAULT: 4.0
shape_matching	Method to adapt student outputs to teacher shape. "projection": Learn a linear projection layer to match dimensions (recommended). "padding": Zero-pad the smaller tensor to match dimensions. "truncation": Truncate the larger tensor to match dimensions. TYPE: str, optional, by default "projection" DEFAULT: 'projection'
reduce	Reduction method. TYPE: str, optional, by default "mean" DEFAULT: 'mean'
sign	Sign multiplier for loss. TYPE: float, optional, by default 1.0 DEFAULT: 1.0

Examples:

>>> loss_fn = DistillationLoss(alpha=0.7, temperature=4.0)
>>> T = torch.randn(100, 10)
>>> S = torch.randn(100, 10)
>>> target = torch.randn(100, 10)
>>> context = {"T": T, "S": S, "target": target}
>>> loss = loss_fn(context)

METHOD	DESCRIPTION
forward	Compute combined distillation loss and task loss (MSE by default).

Source code in spectre/loss/distillation.py

def __init__(
    self,
    alpha: float = 0.7,
    temperature: float = 4.0,
    shape_matching: Literal["projection", "padding", "truncation"] = "projection",
    reduce: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
    context_prefix: str | None = None,
):
    super().__init__(reduce=reduce, sign=sign, context_prefix=context_prefix)

    if not isinstance(alpha, float):
        raise TypeError(f"`alpha` must be a float, got {type(alpha).__name__}.")
    check_in_interval(alpha, "[0, 1]")
    self.alpha = alpha

    if not isinstance(temperature, float):
        raise TypeError(
            f"`temperature` must be a float, got {type(temperature).__name__}."
        )
    check_in_interval(temperature, "(0, inf)")
    self.temperature = temperature

    if shape_matching not in ["projection", "padding", "truncation"]:
        raise ValueError(
            f"Unknown shape_matching method: {shape_matching}. "
            "Choose from 'projection', 'padding', 'truncation'."
        )
    self.shape_matching = shape_matching

Functions

forward(context: dict, **kwargs) -> torch.Tensor

Compute combined distillation loss and task loss (MSE by default).

PARAMETER	DESCRIPTION
context	Context dictionary containing teacher outputs (T), student outputs (S), target (target), and weights (weights). TYPE: dict
**kwargs	Additional keyword arguments. TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Combined loss.

Source code in spectre/loss/distillation.py

def forward(self, context: dict, **kwargs) -> torch.Tensor:
    """
    Compute combined distillation loss and task loss (MSE by default).

    Parameters
    ----------
    context : dict
        Context dictionary containing teacher outputs (`T`), student outputs (`S`),
        target (`target`), and weights (`weights`).

    **kwargs : dict
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        Combined loss.
    """
    T = self._get_context(context, "T", None)
    S = self._get_context(context, "S", None)
    target = self._get_context(context, "target", None)
    weights = self._get_context(context, "weights", None)

    if T is None or S is None:
        raise ValueError(
            "`DistillationLoss` requires `T` (teacher outputs) and "
            "`S` (student outputs) to be provided in the context."
        )

    return distillation_loss(
        S,
        T,
        weights=weights,
        target=target,
        alpha=self.alpha,
        temperature=self.temperature,
        shape_matching=self.shape_matching,
        reduce=self.reduce,
        sign=self.sign,
    )

FeatureMatchingLoss(feature_weight: float = 1.0, shape_matching: Literal['projection', 'padding', 'truncation'] = 'projection', reduce: Literal['sum', 'mean'] = 'mean', sign: float = 1.0, context_prefix: str | None = None)

Bases: Loss

Loss for matching intermediate feature representations.

Computes MSE loss between teacher and student feature maps, automatically handling dimension mismatches through projection.

Loss is defined as \(\frac{1}{L} \sum_{i=1}^L \| F_S^{(i)} - F_T^{(i)} \|^2_2\), where \(L\) is the number of feature layers, and \(F_S^{(i)}\) and \(F_T^{(i)}\) are the student and teacher feature maps at layer \(i\).

PARAMETER	DESCRIPTION
feature_weight	Weight for feature matching loss. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
shape_matching	Method to adapt student outputs to teacher shape. "projection": Learn a linear projection layer to match dimensions. "padding": Zero-pad the smaller tensor to match dimensions. "truncation": Truncate the larger tensor to match dimensions. TYPE: str, optional, by default "projection" DEFAULT: 'projection'
reduce	Reduction method. TYPE: str, optional, by default "mean" DEFAULT: 'mean'
sign	Sign multiplier for loss. TYPE: float, optional, by default 1.0 DEFAULT: 1.0

Examples:

>>> loss_fn = FeatureMatchingLoss(feature_weight=1.0)
>>> T = {"layer1": torch.randn(100, 64), "layer2": torch.randn(100, 128, 8, 8)}
>>> S = {"layer1": torch.randn(100, 32), "layer2": torch.randn(100, 64, 8, 8)}
>>> context = {"T": T, "S": S}
>>> loss = loss_fn(context)

METHOD	DESCRIPTION
forward	Compute feature matching loss between teacher and student features.

Source code in spectre/loss/distillation.py

def __init__(
    self,
    feature_weight: float = 1.0,
    shape_matching: Literal["projection", "padding", "truncation"] = "projection",
    reduce: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
    context_prefix: str | None = None,
):
    super().__init__(reduce=reduce, sign=sign, context_prefix=context_prefix)

    if not isinstance(feature_weight, float):
        raise TypeError(
            f"`feature_weight` must be a float, got {type(feature_weight).__name__}."
        )
    check_in_interval(feature_weight, "[0, inf)")
    self.feature_weight = feature_weight

    if shape_matching not in ["projection", "padding", "truncation"]:
        raise ValueError(
            f"Unknown shape_matching method: {shape_matching}. "
            "Choose from 'projection', 'padding', 'truncation'."
        )
    self.shape_matching = shape_matching

    # Cache for dimension adaptation layers.
    self.projections = {}

Functions

forward(context: dict, **kwargs) -> torch.Tensor

Compute feature matching loss between teacher and student features.

PARAMETER	DESCRIPTION
context	Dictionary containing: "T": Dict of teacher feature tensors. "S": Dict of student feature tensors. "weights": Sample weights (optional). TYPE: dict
**kwargs	Additional keyword arguments. DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Feature matching loss.

Source code in spectre/loss/distillation.py

def forward(
    self,
    context: dict,
    **kwargs,
) -> torch.Tensor:
    """
    Compute feature matching loss between teacher and student features.

    Parameters
    ----------
    context : dict
        Dictionary containing:

        - "T": Dict of teacher feature tensors.
        - "S": Dict of student feature tensors.
        - "weights": Sample weights (optional).

    **kwargs:
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        Feature matching loss.
    """
    T = self._get_context(context, "T", None)
    S = self._get_context(context, "S", None)
    weights = self._get_context(context, "weights", None)

    if T is None or S is None:
        raise ValueError(
            "`FeatureMatchingLoss` requires `T` (teacher features) and "
            "`S` (student features) to be provided in the context."
        )

    return feature_matching_loss(
        S,
        T,
        weights=weights,
        feature_weight=self.feature_weight,
        shape_matching=self.shape_matching,
        reduce=self.reduce,
        sign=self.sign,
    )

SpectralPreservationLoss(kernel_fn: Kernel | str = 'gaussian', kernel_kwargs: dict | None = None, distance_fn: PairwiseDistance | str = 'euclidean', distance_kwargs: dict | None = None, eigenvalue_weight: float = 0.5, eigengap_weight: float = 0.5, reduce: Literal['mean', 'sum'] = 'mean', sign: float = 1.0, symmetric_eigendecomposition: bool = True, context_prefix: str | None = None)

Bases: Loss

Loss for preserving spectral properties in distillation.

Preserves eigenvalue structures, spectral gaps, and kernel properties when distilling from spectral models.

Spectral preservation loss is given as \(w_{\lambda} L_{\lambda} + w_{\Delta\lambda} L_{\Delta\lambda}\), where \(w_{\lambda}\) and \(w_{\Delta\lambda}\) are weights for eigenvalue matching and spectral gap preservation, respectively.

Eigenvalue loss is computed as MSE between teacher and student eigenvalues.

Spectral gap loss is computed as MSE between differences of consecutive eigenvalues (gaps) of teacher and student.

PARAMETER	DESCRIPTION
kernel_fn	Kernel function specification. Kernel object: Pre-configured kernel instance str: Kernel type name (see available options below) Available string options: "gaussian": GaussianKernel "t": Student's TKernel TYPE: Kernel | str, optional, by default "gaussian" DEFAULT: 'gaussian'
kernel_kwargs	Keyword arguments for kernel instantiation when kernel_fn is a string. See kernels for available options. TYPE: dict | None, optional, by default None DEFAULT: None
distance_fn	Distance metric specification. PairwiseDistance object: Pre-configured distance instance str: Distance type name (see available options below) Available string options: "euclidean": Euclidean L2 distance "mahalanobis": Mahalanobis distance "covariance": Covariance-based distance TYPE: PairwiseDistance | str, optional, by default "euclidean" DEFAULT: 'euclidean'
distance_kwargs	Keyword arguments for distance instantiation when distance_fn is a string. TYPE: dict | None, optional, by default None DEFAULT: None
eigenvalue_weight	Weight for eigenvalue matching loss. If 0, eigenvalue loss is not computed. TYPE: float, optional, by default 0.5 DEFAULT: 0.5
eigengap_weight	Weight for spectral gap preservation. If 0, gap loss is not computed. TYPE: float, optional, by default 0.5 DEFAULT: 0.5
reduce	Reduction method. Specifies how the loss is aggregated across samples. "mean": Average loss over all samples. "sum": Sum of losses over all samples. TYPE: Literal["mean", "sum"], optional, by default "mean" DEFAULT: 'mean'
sign	Sign multiplier for loss. Positive for minimization, negative for maximization. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
symmetric_eigendecomposition	Whether to use symmetric eigendecomposition for eigenvalue matching. TYPE: bool, optional, by default True DEFAULT: True

Examples:

>>> # With pre-computed eigenvalues
>>> loss_fn = SpectralPreservationLoss(eigenvalue_weight=1.0, eigengap_weight=0.5)
>>> context = {"T_eigenvalues": t_eigvals, "S_eigenvalues": s_eigvals}
>>> loss = loss_fn(context)
>>>
>>> # With custom distance/kernel functions
>>> from spectre.pairwise_distance import pairwise_distance_mahalanobis
>>> from spectre.kernel import t_kernel
>>> loss_fn = SpectralPreservationLoss(
...     distance_fn=pairwise_distance_mahalanobis,
...     kernel_fn=t_kernel,
...     kernel_kwargs={"alpha": 1.0, "beta": 3.0},
... )
>>> context = {"S": student_output, "T": teacher_output}
>>> loss = loss_fn(context)

METHOD	DESCRIPTION
forward	Compute spectral preservation loss.

Source code in spectre/loss/distillation.py

def __init__(
    self,
    kernel_fn: Kernel | str = "gaussian",
    kernel_kwargs: dict | None = None,
    distance_fn: PairwiseDistance | str = "euclidean",
    distance_kwargs: dict | None = None,
    eigenvalue_weight: float = 0.5,
    eigengap_weight: float = 0.5,
    reduce: Literal["mean", "sum"] = "mean",
    sign: float = 1.0,
    symmetric_eigendecomposition: bool = True,
    context_prefix: str | None = None,
):
    super().__init__(reduce=reduce, sign=sign, context_prefix=context_prefix)

    self.kernel_fn = initialize_kernel_fn(kernel_fn, kernel_kwargs or {})
    self.distance_fn = initialize_distance_fn(distance_fn, distance_kwargs or {})

    if not isinstance(eigenvalue_weight, float):
        raise TypeError(
            f"`eigenvalue_weight` must be a float, got {type(eigenvalue_weight)}."
        )
    check_in_interval(eigenvalue_weight, "[0, 1]")
    self.eigenvalue_weight = eigenvalue_weight

    if not isinstance(eigengap_weight, float):
        raise TypeError(
            f"`eigengap_weight` must be a float, got {type(eigengap_weight)}."
        )
    check_in_interval(eigengap_weight, "[0, 1]")
    self.eigengap_weight = eigengap_weight

    self.symmetric_eigendecomposition = symmetric_eigendecomposition

Functions

forward(context: dict, **kwargs) -> torch.Tensor

Compute spectral preservation loss.

PARAMETER	DESCRIPTION
context	Dictionary containing either eigenvalues or embeddings: Eigenvalues: "T_eigenvalues": Tensor of teacher eigenvalues. "S_eigenvalues": Tensor of student eigenvalues. Embeddings: "S": Student output tensor. "T": Teacher output tensor. Optional: "weights": Sample weights. TYPE: dict
**kwargs	Additional keyword arguments. DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Spectral preservation loss.

Source code in spectre/loss/distillation.py

def forward(
    self,
    context: dict,
    **kwargs,
) -> torch.Tensor:
    """
    Compute spectral preservation loss.

    Parameters
    ----------
    context : dict
        Dictionary containing either eigenvalues or embeddings:

        - Eigenvalues:
            - "T_eigenvalues": Tensor of teacher eigenvalues.
            - "S_eigenvalues": Tensor of student eigenvalues.

        - Embeddings:
            - "S": Student output tensor.
            - "T": Teacher output tensor.

        - Optional:
            - "weights": Sample weights.

    **kwargs:
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        Spectral preservation loss.
    """
    T_eigenvalues = self._get_context(context, "T_eigenvalues", None)
    S_eigenvalues = self._get_context(context, "S_eigenvalues", None)

    S = self._get_context(context, "S", None)
    T = self._get_context(context, "T", None)

    weights = self._get_context(context, "weights", None)

    params = {
        "distance_fn": self.distance_fn,
        "kernel_fn": self.kernel_fn,
        "weights": weights,
        "eigenvalue_weight": self.eigenvalue_weight,
        "eigengap_weight": self.eigengap_weight,
        "reduce": self.reduce,
        "sign": self.sign,
        "symmetric_eigendecomposition": self.symmetric_eigendecomposition,
    }

    # Check first if eigenvalues are provided directly
    if T_eigenvalues is not None and S_eigenvalues is not None:
        return spectral_preservation_loss(
            S_eigenvalues=S_eigenvalues, T_eigenvalues=T_eigenvalues, **params
        )
    # Otherwise, compute eigenvalues from embeddings
    elif S is not None or T is not None:
        return spectral_preservation_loss(S=S, T=T, **params)
    else:
        raise ValueError(
            "`SpectralPreservationLoss` requires either (`T_eigenvalues`, `S_eigenvalues`) "
            "or (`S`, `T`) to be provided in the context."
        )

Functions

distillation_loss(S: torch.Tensor, T: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None, alpha: float = 0.7, temperature: float = 4.0, shape_matching: Literal['projection', 'padding', 'truncation'] = 'projection', reduce: Literal['sum', 'mean'] = 'mean', sign: float = 1.0)

Compute combined distillation loss and task loss (MSE by default).

Combines distillation loss (student vs teacher) with task loss (student vs ground truth) with using a weighted combination.

Loss is given as \(\alpha L_D + (1 - \alpha) L_T\), where \(\alpha\) balances the two components.

Distillation loss is given as \(L_D = t^2 \mathrm{KL}(\text{softmax}(S / t)~\|~\text{softmax}(T / t))\), where \(t\) is the temperature parameter.

Task loss is given as mean squared error \(L_T = \text{MSE}(S, Y)\), where \(S\) is the student output and \(Y\) is the ground truth.

PARAMETER	DESCRIPTION
S	Student model outputs of shape (n_samples, k_features). TYPE: Tensor
T	Teacher model outputs of shape (n_samples, l_features). TYPE: Tensor
weights	Sample weights of shape (n_samples,). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
target	Ground truth targets of shape (n_samples,). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
alpha	Weight for distillation loss vs task loss. alpha=1.0: Pure distillation (no task loss) alpha=0.0: Pure task loss (no distillation) 0 < alpha < 1: Weighted combination TYPE: float, optional, by default 0.7 DEFAULT: 0.7
temperature	Temperature for soft target scaling. TYPE: float, optional, by default 4.0 DEFAULT: 4.0
shape_matching	Method to adapt student outputs to teacher shape. "projection": Learn a linear projection layer to match dimensions. "padding": Zero-pad the smaller tensor to match dimensions. "truncation": Truncate the larger tensor to match dimensions. TYPE: str, optional, by default "projection" DEFAULT: 'projection'
reduce	Reduction method. TYPE: str, optional, by default "mean" DEFAULT: 'mean'
sign	Sign multiplier for loss. TYPE: float, optional, by default 1.0 DEFAULT: 1.0

RETURNS	DESCRIPTION
Tensor	Distillation loss.

Examples:

>>> loss = distillation_loss(S, T, alpha=0.7, temperature=4.0)

Source code in spectre/loss/distillation.py

def distillation_loss(
    S: torch.Tensor,
    T: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
    alpha: float = 0.7,
    temperature: float = 4.0,
    shape_matching: Literal["projection", "padding", "truncation"] = "projection",
    reduce: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
):
    """
    Compute combined distillation loss and task loss (MSE by default).

    Combines distillation loss (student vs teacher) with task loss
    (student vs ground truth) with using a weighted combination.

    Loss is given as $\\alpha L_D + (1 - \\alpha) L_T$, where $\\alpha$ balances
    the two components.

    Distillation loss is given as
    $L_D = t^2 \\mathrm{KL}(\\text{softmax}(S / t)~\\|~\\text{softmax}(T / t))$,
    where $t$ is the temperature parameter.

    Task loss is given as mean squared error
    $L_T = \\text{MSE}(S, Y)$,
    where $S$ is the student output and $Y$ is the ground truth.

    Parameters
    ----------
    S : torch.Tensor
        Student model outputs of shape (n_samples, k_features).

    T : torch.Tensor
        Teacher model outputs of shape (n_samples, l_features).

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape (n_samples,).

    target : torch.Tensor | None, optional, by default None
        Ground truth targets of shape (n_samples,).

    alpha : float, optional, by default 0.7
        Weight for distillation loss vs task loss.

        - `alpha=1.0`: Pure distillation (no task loss)
        - `alpha=0.0`: Pure task loss (no distillation)
        - `0 < alpha < 1`: Weighted combination

    temperature : float, optional, by default 4.0
        Temperature for soft target scaling.

    shape_matching : str, optional, by default "projection"
        Method to adapt student outputs to teacher shape.

        - "projection": Learn a linear projection layer to match dimensions.
        - "padding": Zero-pad the smaller tensor to match dimensions.
        - "truncation": Truncate the larger tensor to match dimensions.

    reduce : str, optional, by default "mean"
        Reduction method.

    sign : float, optional, by default 1.0
        Sign multiplier for loss.

    Returns
    -------
    torch.Tensor
        Distillation loss.

    Examples
    --------
    >>> loss = distillation_loss(S, T, alpha=0.7, temperature=4.0)
    """
    loss = torch.tensor(0.0, device=S.device, dtype=S.dtype)

    # Distillation loss (student learns from teacher).
    if alpha > 0:
        # Adapt student outputs to match teacher output shape.
        S_matched = match_shape(S, T, method=shape_matching)

        # Temperature-scaled soft targets.
        teacher_soft = torch.nn.functional.softmax(T / temperature, dim=-1)
        student_log_soft = torch.nn.functional.log_softmax(
            S_matched / temperature, dim=-1
        )

        # KL divergence with temperature scaling.
        kl_loss = kl_divergence_loss(
            student_log_soft,
            teacher_soft,
            weights=weights,
            reduce=reduce,
        )
        loss += alpha * kl_loss * temperature**2

    # Task loss (student learns from ground truth).
    if alpha < 1.0 and target is not None:
        task_loss = mse_loss(S, target, weights=weights, reduce=reduce, sign=1.0)
        loss += (1 - alpha) * task_loss

    return sign * loss

feature_matching_loss(S: torch.Tensor, T: torch.Tensor, weights: torch.Tensor | None = None, feature_weight: float = 1.0, shape_matching: Literal['projection', 'padding', 'truncation'] = 'projection', reduce: Literal['sum', 'mean'] = 'mean', sign: float = 1.0)

Compute feature matching loss between teacher and student features.

Computes MSE loss between teacher and student feature maps, automatically handling dimension mismatches through projection.

Loss is defined as \(\frac{1}{L} \sum_{i=1}^L \| F_S^{(i)} - F_T^{(i)} \|^2_2\), where \(L\) is the number of feature layers, and \(F_S^{(i)}\) and \(F_T^{(i)}\) are the student and teacher feature maps at layer \(i\).

PARAMETER	DESCRIPTION
S	Student feature dict with layer names as keys and tensors as values. TYPE: Tensor
T	Teacher feature dict with layer names as keys and tensors as values. TYPE: Tensor
weights	Sample weights of shape (n_samples,). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
feature_weight	Weight for feature matching loss. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
shape_matching	Method to adapt student outputs to teacher shape. "projection": Learn a linear projection layer to match dimensions. "padding": Zero-pad the smaller tensor to match dimensions. "truncation": Truncate the larger tensor to match dimensions. TYPE: str, optional, by default "projection" DEFAULT: 'projection'
reduce	Reduction method. TYPE: str, optional, by default "mean" DEFAULT: 'mean'
sign	Sign multiplier for loss. TYPE: float, optional, by default 1.0 DEFAULT: 1.0

RETURNS	DESCRIPTION
Tensor	Feature matching loss.

Examples:

>>> loss = feature_matching_loss(S, T, weights=weights, feature_weight=1.0)

Source code in spectre/loss/distillation.py

def feature_matching_loss(
    S: torch.Tensor,
    T: torch.Tensor,
    weights: torch.Tensor | None = None,
    feature_weight: float = 1.0,
    shape_matching: Literal["projection", "padding", "truncation"] = "projection",
    reduce: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
):
    """
    Compute feature matching loss between teacher and student features.

    Computes MSE loss between teacher and student feature maps, automatically handling
    dimension mismatches through projection.

    Loss is defined as
    $\\frac{1}{L} \\sum_{i=1}^L \\| F_S^{(i)} - F_T^{(i)} \\|^2_2$, where $L$ is
    the number of feature layers, and $F_S^{(i)}$ and $F_T^{(i)}$ are the student
    and teacher feature maps at layer $i$.

    Parameters
    ----------
    S : torch.Tensor
        Student feature dict with layer names as keys and tensors as values.

    T : torch.Tensor
        Teacher feature dict with layer names as keys and tensors as values.

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape (n_samples,).

    feature_weight : float, optional, by default 1.0
        Weight for feature matching loss.

    shape_matching : str, optional, by default "projection"
        Method to adapt student outputs to teacher shape.

        - "projection": Learn a linear projection layer to match dimensions.
        - "padding": Zero-pad the smaller tensor to match dimensions.
        - "truncation": Truncate the larger tensor to match dimensions.

    reduce : str, optional, by default "mean"
        Reduction method.

    sign : float, optional, by default 1.0
        Sign multiplier for loss.

    Returns
    -------
    torch.Tensor
        Feature matching loss.

    Examples
    --------
    >>> loss = feature_matching_loss(S, T, weights=weights, feature_weight=1.0)
    """
    fm_loss = torch.tensor(0.0)
    n_layers = 0

    for layer_name in T:
        if layer_name in S:
            T_features = T[layer_name]
            S_features = S[layer_name]

            # Adapt student features to teacher dimensions.
            S_features = match_shape(S_features, T_features, method=shape_matching)

            loss = mse_loss(S_features, T_features, weights=weights, reduce=reduce)

            fm_loss += loss
            n_layers += 1

    if n_layers > 0:
        fm_loss = fm_loss / n_layers

    return sign * feature_weight * fm_loss

spectral_preservation_loss(*, S: torch.Tensor | None = None, T: torch.Tensor | None = None, S_eigenvalues: torch.Tensor | None = None, T_eigenvalues: torch.Tensor | None = None, distance_fn: PairwiseDistance, kernel_fn: Kernel, weights: torch.Tensor | None = None, eigenvalue_weight: float = 1.0, eigengap_weight: float = 0.5, reduce: Literal['mean', 'sum'] = 'mean', sign: float = 1.0, symmetric_eigendecomposition: bool = True)

Loss for preserving spectral properties in distillation.

Preserves eigenvalue structures, spectral gaps, and kernel properties when distilling from spectral models.

Spectral preservation loss is given as \(w_{\lambda} L_{\lambda} + w_{\Delta\lambda} L_{\Delta\lambda}\), where \(w_{\lambda}\) and \(w_{\Delta\lambda}\) are weights for eigenvalue matching and spectral gap preservation, respectively.

Eigenvalue loss is computed as MSE between teacher and student eigenvalues.

Spectral gap loss is computed as MSE between differences of consecutive eigenvalues (gaps) of teacher and student.

PARAMETER	DESCRIPTION
S	Student output of shape (n_samples, k_features). Required if eigenvalues not provided. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
T	Teacher output of shape (n_samples, l_features). Required if eigenvalues not provided. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
S_eigenvalues	Pre-computed student eigenvalues. If provided, eigenvalues are ignored. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
T_eigenvalues	Pre-computed teacher eigenvalues. If provided, eigenvalues are ignored. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
distance_fn	Pairwise distance function. TYPE: PairwiseDistance
kernel_fn	Kernel function. TYPE: Kernel
weights	Sample weights of shape (n_samples,). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
eigenvalue_weight	Weight for eigenvalue matching loss. If 0, eigenvalue loss is not computed. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
eigengap_weight	Weight for spectral gap preservation. If 0, gap loss is not computed. TYPE: float, optional, by default 0.5 DEFAULT: 0.5
reduce	Reduction method. Specifies how the loss is aggregated across eigenvalues. "mean": Average loss over all eigenvalues. "sum": Sum of losses over all eigenvalues. TYPE: Literal["mean", "sum"], optional, by default "mean" DEFAULT: 'mean'
sign	Sign multiplier for loss. Positive for minimization, negative for maximization. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
symmetric_eigendecomposition	Whether to use symmetric eigendecomposition for eigenvalue matching. TYPE: bool, optional, by default True DEFAULT: True

RETURNS	DESCRIPTION
Tensor	Spectral preservation loss.

Examples:

>>> # With pre-computed eigenvalues
>>> loss = spectral_preservation_loss(
...     T_eigenvalues=t_eigvals,
...     S_eigenvalues=s_eigvals,
...     eigenvalue_weight=1.0,
...     eigengap_weight=0.5,
... )
>>>
>>> # With embeddings and default distance/kernel
>>> loss = spectral_preservation_loss(
...     S=student_output,
...     T=teacher_output,
...     eigenvalue_weight=1.0,
...     eigengap_weight=0.5,
... )
>>>
>>> # With custom distance/kernel functions
>>> from spectre.pairwise_distance import pairwise_distance_mahalanobis
>>> from spectre.kernel import t_kernel
>>> loss = spectral_preservation_loss(
...     S=student_output,
...     T=teacher_output,
...     distance_fn=pairwise_distance_mahalanobis,
...     kernel_fn=t_kernel,
...     kernel_kwargs={"alpha": 1.0, "beta": 3.0},
... )

Source code in spectre/loss/distillation.py

def spectral_preservation_loss(
    *,
    S: torch.Tensor | None = None,
    T: torch.Tensor | None = None,
    S_eigenvalues: torch.Tensor | None = None,
    T_eigenvalues: torch.Tensor | None = None,
    distance_fn: PairwiseDistance,
    kernel_fn: Kernel,
    weights: torch.Tensor | None = None,
    eigenvalue_weight: float = 1.0,
    eigengap_weight: float = 0.5,
    reduce: Literal["mean", "sum"] = "mean",
    sign: float = 1.0,
    symmetric_eigendecomposition: bool = True,
):
    """
    Loss for preserving spectral properties in distillation.

    Preserves eigenvalue structures, spectral gaps, and kernel properties
    when distilling from spectral models.

    Spectral preservation loss is given as
    $w_{\\lambda} L_{\\lambda} + w_{\\Delta\\lambda} L_{\\Delta\\lambda}$,
    where $w_{\\lambda}$ and $w_{\\Delta\\lambda}$ are weights for
    eigenvalue matching and spectral gap preservation, respectively.

    Eigenvalue loss is computed as MSE between teacher and student eigenvalues.

    Spectral gap loss is computed as MSE between differences of consecutive
    eigenvalues (gaps) of teacher and student.

    Parameters
    ----------
    S : torch.Tensor | None, optional, by default None
        Student output of shape (n_samples, k_features). Required if eigenvalues
        not provided.

    T : torch.Tensor | None, optional, by default None
        Teacher output of shape (n_samples, l_features). Required if eigenvalues
        not provided.

    S_eigenvalues : torch.Tensor | None, optional, by default None
        Pre-computed student eigenvalues. If provided, eigenvalues are ignored.

    T_eigenvalues : torch.Tensor | None, optional, by default None
        Pre-computed teacher eigenvalues. If provided, eigenvalues are ignored.

    distance_fn : PairwiseDistance
        Pairwise distance function.

    kernel_fn : Kernel
        Kernel function.

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape (n_samples,).

    eigenvalue_weight : float, optional, by default 1.0
        Weight for eigenvalue matching loss. If 0, eigenvalue loss is not computed.

    eigengap_weight : float, optional, by default 0.5
        Weight for spectral gap preservation. If 0, gap loss is not computed.

    reduce : Literal["mean", "sum"], optional, by default "mean"
        Reduction method. Specifies how the loss is aggregated across eigenvalues.

        - "mean": Average loss over all eigenvalues.
        - "sum": Sum of losses over all eigenvalues.

    sign : float, optional, by default 1.0
        Sign multiplier for loss. Positive for minimization, negative for maximization.

    symmetric_eigendecomposition : bool, optional, by default True
        Whether to use symmetric eigendecomposition for eigenvalue matching.

    Returns
    -------
    torch.Tensor
        Spectral preservation loss.

    Examples
    --------
    >>> # With pre-computed eigenvalues
    >>> loss = spectral_preservation_loss(
    ...     T_eigenvalues=t_eigvals,
    ...     S_eigenvalues=s_eigvals,
    ...     eigenvalue_weight=1.0,
    ...     eigengap_weight=0.5,
    ... )
    >>>
    >>> # With embeddings and default distance/kernel
    >>> loss = spectral_preservation_loss(
    ...     S=student_output,
    ...     T=teacher_output,
    ...     eigenvalue_weight=1.0,
    ...     eigengap_weight=0.5,
    ... )
    >>>
    >>> # With custom distance/kernel functions
    >>> from spectre.pairwise_distance import pairwise_distance_mahalanobis
    >>> from spectre.kernel import t_kernel
    >>> loss = spectral_preservation_loss(
    ...     S=student_output,
    ...     T=teacher_output,
    ...     distance_fn=pairwise_distance_mahalanobis,
    ...     kernel_fn=t_kernel,
    ...     kernel_kwargs={"alpha": 1.0, "beta": 3.0},
    ... )
    """
    # If eigenvalues are provided, use them directly
    if S_eigenvalues is not None and T_eigenvalues is not None:
        loss = torch.tensor(0.0, device=S_eigenvalues.device, dtype=S_eigenvalues.dtype)

        # Eigenvalue matching loss
        if eigenvalue_weight > 0:
            eigval_loss = mse_loss(S_eigenvalues, T_eigenvalues, reduce=reduce)
            loss += eigenvalue_weight * eigval_loss

        # Spectral gap preservation
        if eigengap_weight > 0:
            teacher_gaps = T_eigenvalues[:-1] - T_eigenvalues[1:]
            student_gaps = S_eigenvalues[:-1] - S_eigenvalues[1:]
            gap_loss = mse_loss(student_gaps, teacher_gaps, reduce=reduce)
            loss += eigengap_weight * gap_loss

        return sign * loss

    # Otherwise, compute eigenvalues from embeddings
    if S is None or T is None:
        raise ValueError(
            "`spectral_preservation_loss` requires either (`S_eigenvalues`, `T_eigenvalues`) "
            "or (`S`, `T`) to be provided."
        )

    # Compute kernel matrices and eigenvalues
    s_kernel = kernel_fn(distance_fn(S), weights=weights)
    t_kernel = kernel_fn(distance_fn(T), weights=weights)

    S_eigenvalues = eigvals(
        s_kernel,
        n_eigen=0,
        symmetric=symmetric_eigendecomposition,
        sort_descending=True,
        UPLO="U",
    )
    T_eigenvalues = eigvals(
        t_kernel,
        n_eigen=0,
        symmetric=symmetric_eigendecomposition,
        sort_descending=True,
        UPLO="U",
    )

    # Recursively call with computed eigenvalues
    return spectral_preservation_loss(
        S_eigenvalues=S_eigenvalues,
        T_eigenvalues=T_eigenvalues,
        distance_fn=distance_fn,
        kernel_fn=kernel_fn,
        weights=weights,
        eigenvalue_weight=eigenvalue_weight,
        eigengap_weight=eigengap_weight,
        reduce=reduce,
        sign=sign,
    )

match_shape(S: torch.Tensor, T: torch.Tensor, method: str = 'projection') -> torch.Tensor

Match dimensions between student and teacher tensors (2D only).

• Learnable adaptation: Creates a linear layer that can be trained
• Preserves information: Projects high-dimensional features to lower dimensions
• Bidirectional: Works whether student is smaller or larger than teacher
• Compatible with backprop: Gradients flow through the projection layer

PARAMETER	DESCRIPTION
S	Student model output (n_samples, n_features). TYPE: Tensor
T	Teacher model output (n_samples, n_features). TYPE: Tensor
method	Method for dimension matching ('projection', 'padding', 'truncation'). 'projection': Learnable linear transformation 'padding': Zero-padding to match dimensions 'truncation': Cropping to match dimensions TYPE: str DEFAULT: 'projection'

RETURNS	DESCRIPTION
Tensor	Student tensor adapted to teacher dimensions.

Examples:

>>> S = torch.randn(100, 32)  # Student output
>>> T = torch.randn(100, 64)  # Teacher output
>>> S_matched = match_shape(S, T, method="projection")  # S adapted to T's shape
>>> S_matched.shape
torch.Size([100, 64])

Source code in spectre/loss/distillation.py

def match_shape(
    S: torch.Tensor,
    T: torch.Tensor,
    method: str = "projection",
) -> torch.Tensor:
    """
    Match dimensions between student and teacher tensors (2D only).

    - Learnable adaptation: Creates a linear layer that can be trained
    - Preserves information: Projects high-dimensional features to lower dimensions
    - Bidirectional: Works whether student is smaller or larger than teacher
    - Compatible with backprop: Gradients flow through the projection layer

    Parameters
    ----------
    S : torch.Tensor
        Student model output (n_samples, n_features).

    T : torch.Tensor
        Teacher model output (n_samples, n_features).

    method : str
        Method for dimension matching ('projection', 'padding', 'truncation').

        - 'projection': Learnable linear transformation
        - 'padding': Zero-padding to match dimensions
        - 'truncation': Cropping to match dimensions

    Returns
    -------
    torch.Tensor
        Student tensor adapted to teacher dimensions.

    Examples
    --------
    >>> S = torch.randn(100, 32)  # Student output
    >>> T = torch.randn(100, 64)  # Teacher output
    >>> S_matched = match_shape(S, T, method="projection")  # S adapted to T's shape
    >>> S_matched.shape
    torch.Size([100, 64])
    """
    if S.shape == T.shape:
        return S

    # Do not support Conv2D or higher dimensions for now
    if S.dim() != 2 or T.dim() != 2:
        return S

    if method == "projection":
        # Linear projection to match last dimension
        student_dim = S.shape[-1]
        teacher_dim = T.shape[-1]

        if student_dim != teacher_dim:
            projection = torch.nn.Linear(student_dim, teacher_dim, bias=False)
            projection = projection.to(S.device)
            return projection(S)

    elif method == "padding":
        # Pad smaller tensor
        if S.shape[-1] < T.shape[-1]:
            pad_size = T.shape[-1] - S.shape[-1]
            padding = torch.zeros(*S.shape[:-1], pad_size, device=S.device)
            return torch.cat([S, padding], dim=-1)

    elif method == "truncation":
        # Truncate larger tensor
        if S.shape[-1] > T.shape[-1]:
            return S[..., : T.shape[-1]]

    return S