features

shrinkai.distillation.losses.features

Feature-based Knowledge Distillation loss functions.

This module provides loss functions designed to operate on the intermediate representations (hidden layers, attention maps, feature maps) of neural networks.

Unlike logit-based distillation which only penalizes the final output, feature distillation forces the student to learn the internal reasoning process and hierarchical representations of the teacher. This is particularly crucial when the teacher is much deeper than the student (e.g., FitNets, TinyBERT).

Available Losses

• FeatureLoss: Aligns standard activations (Conv/Linear) using MSE, L1, or Cosine.
• AttentionMapLoss: Aligns Transformer attention matrices (TinyBERT style).
• GramMatrixLoss: Aligns feature co-occurrences for style/texture transfer.

Note

Feature dimensions often differ between student and teacher models. It is the user's responsibility to apply a projection layer (e.g., a 1x1 Conv or a Linear layer) to the student's features before passing them to these losses. To deal with, you can use a ProjectedFeatureLoss, it is a solution, but not necessarily the best one for your problem.

Examples:

>>> from shrinkai.distillation.losses import FeatureLoss
>>> criterion = FeatureLoss(loss_type="mse", normalize=True)
>>> loss = criterion(student_hidden_states, teacher_hidden_states)

Classes:

Name	Description
AttentionMapLoss	Distillation loss for Transformer attention maps.
FeatureLoss	Computes the loss between intermediate feature maps of the Teacher and Student
GramMatrixLoss	Distillation loss based on Gram Matrices for style and texture transfer

Classes

AttentionMapLoss

Bases: BaseDistillationLoss

Distillation loss for Transformer attention maps.

This loss forces the student model to mimic the attention patterns of the teacher. It operates on attention matrices, typically of shape: [Batch, Num_Heads, Seq_Len, Seq_Len].

Equation

\[ L_{attn}(A_s, A_t) = \begin{cases} \text{MSE}(A_s, A_t) & \text{(mse, Jiao et al. (2020), TinyBERT)} \\ \text{KL}\left(\text{softmax}(A_t) \parallel \text{softmax}(A_s)\right) & \text{(kl, Wang et al. (2020), MiniLM)} \end{cases} \]

TinyBERT applies MSE directly on the unnormalized attention scores (before softmax); MiniLM instead matches the softmax-normalized attention distributions via KL divergence, with the teacher as the reference distribution. Note: DistilBERT (Sanh et al. (2019)) is often mentioned alongside these, but its own distillation objective operates on the output logits and last hidden state, not on attention maps.

Attributes:

Name	Type	Description
loss_type	str	The metric to use ('mse' or 'kl').

Methods:

Name	Description
__init__	Initializes the AttentionMapLoss.
forward	Computes the attention map distillation loss.

Source code in src/shrinkai/distillation/losses/features.py

class AttentionMapLoss(BaseDistillationLoss):
    r"""Distillation loss for Transformer attention maps.

    This loss forces the student model to mimic the attention patterns of the teacher.
    It operates on attention matrices, typically of shape:
    [Batch, Num_Heads, Seq_Len, Seq_Len].

    Equation:
        $$
        L_{attn}(A_s, A_t) =
        \begin{cases}
        \text{MSE}(A_s, A_t) & \text{(mse, Jiao et al. (2020), TinyBERT)} \\
        \text{KL}\left(\text{softmax}(A_t) \parallel \text{softmax}(A_s)\right) & \text{(kl, Wang et al. (2020), MiniLM)}
        \end{cases}
        $$

        TinyBERT applies MSE directly on the unnormalized attention scores (before
        softmax); MiniLM instead matches the softmax-normalized attention
        *distributions* via KL divergence, with the teacher as the reference
        distribution. Note: DistilBERT (Sanh et al. (2019)) is often mentioned
        alongside these, but its own distillation objective operates on the output
        logits and last hidden state, not on attention maps.

    Attributes:
        loss_type (str): The metric to use ('mse' or 'kl').
    """  # noqa: E501

    def __init__(self, loss_type: Literal["mse", "kl"] = "mse") -> None:
        """Initializes the AttentionMapLoss.

        Args:
            loss_type: Distance metric ('mse' or 'kl'). Defaults to 'mse'.
                    If 'kl' is used, inputs must be unnormalized logits (before softmax),
                    and the KL divergence will be applied across the last dimension.

        Raises:
            ValueError: If an unsupported `loss_type` is provided.
        """
        super().__init__()
        if loss_type not in ["mse", "kl"]:
            raise ValueError(f"Unsupported loss_type '{loss_type}'. Use 'mse' or 'kl'.")

        self.loss_type = loss_type

    def _compute_distance(self, s_map: torch.Tensor, t_map: torch.Tensor) -> torch.Tensor:
        """Computes the loss between two attention matrices."""
        if s_map.shape != t_map.shape:
            raise ValueError(
                f"Attention map shape mismatch: Student {s_map.shape} vs Teacher {t_map.shape}. "
                "The number of heads and sequence lengths must match. If your student has "
                "fewer heads, consider aligning specific heads before passing them to this loss."
            )

        if self.loss_type == "mse":
            return F.mse_loss(s_map, t_map)

        elif self.loss_type == "kl":
            log_soft_s = F.log_softmax(s_map, dim=-1)
            soft_t = F.softmax(t_map, dim=-1)
            return F.kl_div(input=log_soft_s, target=soft_t, reduction="batchmean")

        raise ValueError("Invalid loss type.")

    def forward(
        self,
        student_outputs: torch.Tensor | dict[str, torch.Tensor],
        teacher_outputs: torch.Tensor | dict[str, torch.Tensor],
        labels: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Computes the attention map distillation loss.

        Supports comparing single tensors or dictionaries of tensors. If dictionaries
        are provided, it computes the average loss across all matching keys.

        Args:
            student_outputs: Tensor or dict of attention matrices from the student.
            teacher_outputs: Tensor or dict of attention matrices from the teacher.
            labels: Ground-truth labels (ignored, kept for API compatibility).

        Returns:
            torch.Tensor: Aggregated scalar loss value.
        """
        return self._apply_feature_distance(
            student_outputs, teacher_outputs, self._compute_distance
        )

Methods:

__init__

__init__(loss_type: Literal['mse', 'kl'] = 'mse') -> None

Initializes the AttentionMapLoss.

Parameters:

Name	Type	Description	Default
loss_type	Literal['mse', 'kl']	Distance metric ('mse' or 'kl'). Defaults to 'mse'. If 'kl' is used, inputs must be unnormalized logits (before softmax), and the KL divergence will be applied across the last dimension.	'mse'

Raises:

Type	Description
ValueError	If an unsupported loss_type is provided.

Source code in src/shrinkai/distillation/losses/features.py

def __init__(self, loss_type: Literal["mse", "kl"] = "mse") -> None:
    """Initializes the AttentionMapLoss.

    Args:
        loss_type: Distance metric ('mse' or 'kl'). Defaults to 'mse'.
                If 'kl' is used, inputs must be unnormalized logits (before softmax),
                and the KL divergence will be applied across the last dimension.

    Raises:
        ValueError: If an unsupported `loss_type` is provided.
    """
    super().__init__()
    if loss_type not in ["mse", "kl"]:
        raise ValueError(f"Unsupported loss_type '{loss_type}'. Use 'mse' or 'kl'.")

    self.loss_type = loss_type

forward

forward(
    student_outputs: Tensor | dict[str, Tensor],
    teacher_outputs: Tensor | dict[str, Tensor],
    labels: Tensor | None = None,
) -> torch.Tensor

Computes the attention map distillation loss.

Supports comparing single tensors or dictionaries of tensors. If dictionaries are provided, it computes the average loss across all matching keys.

Parameters:

Name	Type	Description	Default
student_outputs	Tensor | dict[str, Tensor]	Tensor or dict of attention matrices from the student.	required
teacher_outputs	Tensor | dict[str, Tensor]	Tensor or dict of attention matrices from the teacher.	required
labels	Tensor | None	Ground-truth labels (ignored, kept for API compatibility).	None

Returns:

Type	Description
Tensor	torch.Tensor: Aggregated scalar loss value.

Source code in src/shrinkai/distillation/losses/features.py

def forward(
    self,
    student_outputs: torch.Tensor | dict[str, torch.Tensor],
    teacher_outputs: torch.Tensor | dict[str, torch.Tensor],
    labels: torch.Tensor | None = None,
) -> torch.Tensor:
    """Computes the attention map distillation loss.

    Supports comparing single tensors or dictionaries of tensors. If dictionaries
    are provided, it computes the average loss across all matching keys.

    Args:
        student_outputs: Tensor or dict of attention matrices from the student.
        teacher_outputs: Tensor or dict of attention matrices from the teacher.
        labels: Ground-truth labels (ignored, kept for API compatibility).

    Returns:
        torch.Tensor: Aggregated scalar loss value.
    """
    return self._apply_feature_distance(
        student_outputs, teacher_outputs, self._compute_distance
    )

FeatureLoss

Bases: BaseDistillationLoss

Computes the loss between intermediate feature maps of the Teacher and Student (Romero et al. (2015), FitNets).

This loss encourages the Student to mimic the internal representations (activations) of the Teacher. It assumes that the spatial and channel dimensions of the compared features have already been matched.

Equation

\[L_{feature}\left(\hat{f}_s, \hat{f}_t\right) = d\left(\hat{f}_s, \hat{f}_t\right), \quad d \in \{\text{MSE}, \ \text{L1}, \ 1 - \cos\}\]

where \(\hat{f}_s, \hat{f}_t\) are the (optionally L2-normalized) flattened student and teacher feature tensors. The original FitNets "hint" loss corresponds to the MSE case; L1 and cosine are natural extensions of the same idea.

Attributes:

Name	Type	Description
loss_type	str	Type of loss to compute ('mse', 'l1', or 'cosine').
normalize	bool	If True, L2-normalizes the feature vectors before computing the loss. This is often useful to match the "direction" of features regardless of their magnitude.

Methods:

Name	Description
__init__	Initializes the FeatureLoss.
forward	Computes the feature distillation loss.

Source code in src/shrinkai/distillation/losses/features.py

class FeatureLoss(BaseDistillationLoss):
    r"""Computes the loss between intermediate feature maps of the Teacher and Student
    (Romero et al. (2015), FitNets).

    This loss encourages the Student to mimic the internal representations (activations)
    of the Teacher. It assumes that the spatial and channel dimensions of the compared
    features have already been matched.

    Equation:
        $$L_{feature}\left(\hat{f}_s, \hat{f}_t\right) = d\left(\hat{f}_s, \hat{f}_t\right), \quad d \in \{\text{MSE}, \ \text{L1}, \ 1 - \cos\}$$

        where $\hat{f}_s, \hat{f}_t$ are the (optionally L2-normalized) flattened student
        and teacher feature tensors. The original FitNets "hint" loss corresponds to the
        MSE case; L1 and cosine are natural extensions of the same idea.

    Attributes:
        loss_type (str): Type of loss to compute ('mse', 'l1', or 'cosine').
        normalize (bool): If True, L2-normalizes the feature vectors before computing
            the loss. This is often useful to match the "direction" of features
            regardless of their magnitude.
    """  # noqa: E501

    def __init__(
        self,
        loss_type: Literal["mse", "l1", "cosine"] = "mse",
        normalize: bool = False,
    ) -> None:
        """Initializes the FeatureLoss.

        Args:
            loss_type: Distance metric to use ('mse', 'l1', or 'cosine'). Defaults to 'mse'.
            normalize: Whether to apply L2 normalization to features before comparison.
                Defaults to False.

        Raises:
            ValueError: If an unsupported `loss_type` is provided.
        """
        super().__init__()
        if loss_type not in ["mse", "l1", "cosine"]:
            raise ValueError(f"Unsupported loss_type '{loss_type}'. Use 'mse', 'l1', or 'cosine'.")

        self.loss_type = loss_type
        self.normalize = normalize

    def _compute_distance(
        self, student_feat: torch.Tensor, teacher_feat: torch.Tensor
    ) -> torch.Tensor:
        """Computes the specified distance metric between two feature tensors."""
        if student_feat.shape != teacher_feat.shape:
            raise ValueError(
                f"Feature shape mismatch: Student {student_feat.shape} vs"
                f"Teacher {teacher_feat.shape}."
                "Ensure dimensions match, or apply a projection layer"
                "to the student features before computing the loss."
            )

        if self.normalize:
            student_feat = F.normalize(student_feat.view(student_feat.size(0), -1), p=2, dim=-1)
            teacher_feat = F.normalize(teacher_feat.view(teacher_feat.size(0), -1), p=2, dim=-1)

        if self.loss_type == "mse":
            return F.mse_loss(student_feat, teacher_feat)
        if self.loss_type == "l1":
            return F.l1_loss(student_feat, teacher_feat)
        if self.loss_type == "cosine":
            s_flat = student_feat.view(student_feat.size(0), -1)
            t_flat = teacher_feat.view(teacher_feat.size(0), -1)
            target = torch.ones(s_flat.size(0), device=s_flat.device)
            return F.cosine_embedding_loss(s_flat, t_flat, target)

        raise ValueError("Invalid loss type.")

    def forward(
        self,
        student_outputs: torch.Tensor | dict[str, torch.Tensor],
        teacher_outputs: torch.Tensor | dict[str, torch.Tensor],
        labels: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Computes the feature distillation loss.

        Supports comparing single tensors or dictionaries of tensors. If dictionaries
        are provided, it computes the average loss across all matching keys.

        Args:
            student_outputs: Tensor or dict of intermediate feature tensors from the student.
            teacher_outputs: Tensor or dict of intermediate feature tensors from the teacher.
            labels: Ground-truth labels (ignored, kept for API compatibility).

        Returns:
            torch.Tensor: Aggregated scalar loss value.

        Raises:
            TypeError: If input types for student and teacher do not match.
            ValueError: If dict keys do not match between student and teacher.
        """
        return self._apply_feature_distance(
            student_outputs, teacher_outputs, self._compute_distance
        )

Methods:

__init__

__init__(
    loss_type: Literal["mse", "l1", "cosine"] = "mse",
    normalize: bool = False,
) -> None

Initializes the FeatureLoss.

Parameters:

Name	Type	Description	Default
loss_type	Literal['mse', 'l1', 'cosine']	Distance metric to use ('mse', 'l1', or 'cosine'). Defaults to 'mse'.	'mse'
normalize	bool	Whether to apply L2 normalization to features before comparison. Defaults to False.	False

Raises:

Type	Description
ValueError	If an unsupported loss_type is provided.

Source code in src/shrinkai/distillation/losses/features.py

def __init__(
    self,
    loss_type: Literal["mse", "l1", "cosine"] = "mse",
    normalize: bool = False,
) -> None:
    """Initializes the FeatureLoss.

    Args:
        loss_type: Distance metric to use ('mse', 'l1', or 'cosine'). Defaults to 'mse'.
        normalize: Whether to apply L2 normalization to features before comparison.
            Defaults to False.

    Raises:
        ValueError: If an unsupported `loss_type` is provided.
    """
    super().__init__()
    if loss_type not in ["mse", "l1", "cosine"]:
        raise ValueError(f"Unsupported loss_type '{loss_type}'. Use 'mse', 'l1', or 'cosine'.")

    self.loss_type = loss_type
    self.normalize = normalize

forward

forward(
    student_outputs: Tensor | dict[str, Tensor],
    teacher_outputs: Tensor | dict[str, Tensor],
    labels: Tensor | None = None,
) -> torch.Tensor

Computes the feature distillation loss.

Supports comparing single tensors or dictionaries of tensors. If dictionaries are provided, it computes the average loss across all matching keys.

Parameters:

Name	Type	Description	Default
student_outputs	Tensor | dict[str, Tensor]	Tensor or dict of intermediate feature tensors from the student.	required
teacher_outputs	Tensor | dict[str, Tensor]	Tensor or dict of intermediate feature tensors from the teacher.	required
labels	Tensor | None	Ground-truth labels (ignored, kept for API compatibility).	None

Returns:

Type	Description
Tensor	torch.Tensor: Aggregated scalar loss value.

Raises:

Type	Description
TypeError	If input types for student and teacher do not match.
ValueError	If dict keys do not match between student and teacher.

Source code in src/shrinkai/distillation/losses/features.py

def forward(
    self,
    student_outputs: torch.Tensor | dict[str, torch.Tensor],
    teacher_outputs: torch.Tensor | dict[str, torch.Tensor],
    labels: torch.Tensor | None = None,
) -> torch.Tensor:
    """Computes the feature distillation loss.

    Supports comparing single tensors or dictionaries of tensors. If dictionaries
    are provided, it computes the average loss across all matching keys.

    Args:
        student_outputs: Tensor or dict of intermediate feature tensors from the student.
        teacher_outputs: Tensor or dict of intermediate feature tensors from the teacher.
        labels: Ground-truth labels (ignored, kept for API compatibility).

    Returns:
        torch.Tensor: Aggregated scalar loss value.

    Raises:
        TypeError: If input types for student and teacher do not match.
        ValueError: If dict keys do not match between student and teacher.
    """
    return self._apply_feature_distance(
        student_outputs, teacher_outputs, self._compute_distance
    )

GramMatrixLoss

Bases: BaseDistillationLoss

Distillation loss based on Gram Matrices for style and texture transfer (Gatys et al. (2016)).

Instead of forcing the student to match the exact spatial activations of the teacher (which is strict and requires identical spatial dimensions), this loss forces the student to match the channel-wise feature correlations (co-occurrence).

This is highly effective for Generative tasks, Super-Resolution, or making a student network mimic the global "texture" representation of a teacher.

Equation

\[ \begin{aligned} G_{ij} = \frac{1}{C \cdot N} \sum_{k=1}^{N} F_{ik} F_{jk} \\ L_{gram} = d\left(G_s, G_t\right) \end{aligned} \]

where \(F \in \mathbb{R}^{C \times N}\) is a feature map flattened over its \(C\) channels and \(N\) spatial (or temporal) locations, \(G \in \mathbb{R}^{C \times C}\) is its Gram matrix of channel-wise correlations, and \(d\) is the configured distance (MSE, L1, or cosine).

Attributes:

Name	Type	Description
loss_type	str	The distance metric to apply on the Gram matrices ('mse' or 'l1').

Methods:

Name	Description
__init__	Initializes the GramMatrixLoss.
forward	Computes the Gram matrix distillation loss.

Source code in src/shrinkai/distillation/losses/features.py

class GramMatrixLoss(BaseDistillationLoss):
    r"""Distillation loss based on Gram Matrices for style and texture transfer
    (Gatys et al. (2016)).

    Instead of forcing the student to match the exact spatial activations of the
    teacher (which is strict and requires identical spatial dimensions), this loss
    forces the student to match the channel-wise feature correlations (co-occurrence).

    This is highly effective for Generative tasks, Super-Resolution, or making a
    student network mimic the global "texture" representation of a teacher.

    Equation:
        $$
        \begin{aligned}
        G_{ij} = \frac{1}{C \cdot N} \sum_{k=1}^{N} F_{ik} F_{jk} \\
        L_{gram} = d\left(G_s, G_t\right)
        \end{aligned}
        $$

        where $F \in \mathbb{R}^{C \times N}$ is a feature map flattened over its
        $C$ channels and $N$ spatial (or temporal) locations, $G \in \mathbb{R}^{C
        \times C}$ is its Gram matrix of channel-wise correlations, and $d$ is the
        configured distance (MSE, L1, or cosine).

    Attributes:
        loss_type (str): The distance metric to apply on the Gram matrices ('mse' or 'l1').
    """  # noqa: E501

    def __init__(self, loss_type: Literal["mse", "l1", "cosine"] = "mse") -> None:
        """Initializes the GramMatrixLoss.

        Args:
            loss_type: Distance metric ('mse' or 'l1'). Defaults to 'mse'.

        Raises:
            ValueError: If an unsupported `loss_type` is provided.
        """
        super().__init__()
        if loss_type not in ["mse", "l1", "cosine"]:
            raise ValueError(f"Unsupported loss_type '{loss_type}'. Use 'mse', 'l1' or 'cosine'.")

        self.loss_type = loss_type

    def _compute_gram_matrix(self, x: torch.Tensor) -> torch.Tensor:
        """Computes the normalized Gram matrix of a feature tensor.

        Args:
            x: Input tensor of shape [B, C, H, W] or [B, C, L].

        Returns:
            torch.Tensor: Gram matrix of shape [B, C, C].
        """
        if x.dim() < 3:
            raise ValueError(
                f"Gram matrix requires at least 3D tensors (Batch, Channels, Spatial/Temporal). "
                f"Got tensor of shape {x.shape} (dim={x.dim()})."
            )
        batch_size, channels = x.size(0), x.size(1)

        # [B, C, H, W] -> [B, C, H*W]
        x_flat = x.view(batch_size, channels, -1)
        num_elements = x_flat.size(2)

        gram = torch.bmm(x_flat, x_flat.transpose(1, 2))
        gram = gram / (channels * num_elements)

        return gram

    def _compute_distance(self, s_feat: torch.Tensor, t_feat: torch.Tensor) -> torch.Tensor:
        """Computes the distance between the Gram matrices of student and teacher."""
        if s_feat.size(1) != t_feat.size(1):
            raise ValueError(
                f"Channel dimension mismatch: Student has {s_feat.size(1)} channels, "
                f"Teacher has {t_feat.size(1)} channels. The Gram matrix requires identical "
                f"channel dimensions. Apply a 1x1 Conv (projector) to the student features first."
            )

        s_gram = self._compute_gram_matrix(s_feat)
        t_gram = self._compute_gram_matrix(t_feat)

        if self.loss_type == "mse":
            return F.mse_loss(s_gram, t_gram)
        elif self.loss_type == "l1":
            return F.l1_loss(s_gram, t_gram)
        elif self.loss_type == "cosine":
            s_flat = s_gram.view(s_gram.size(0), -1)
            t_flat = t_gram.view(t_gram.size(0), -1)
            target = torch.ones(s_flat.size(0), device=s_flat.device)
            return F.cosine_embedding_loss(s_flat, t_flat, target)

        raise ValueError("Invalid loss type.")

    def forward(
        self,
        student_outputs: torch.Tensor | dict[str, torch.Tensor],
        teacher_outputs: torch.Tensor | dict[str, torch.Tensor],
        labels: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Computes the Gram matrix distillation loss.

        Args:
            student_outputs: Tensor or dict of intermediate features [B, C, H, W].
            teacher_outputs: Tensor or dict of intermediate features [B, C, H, W].
            labels: Ground-truth labels (ignored, kept for API compatibility).

        Returns:
            torch.Tensor: Aggregated scalar loss value.
        """
        return self._apply_feature_distance(
            student_outputs, teacher_outputs, self._compute_distance
        )

Methods:

__init__

__init__(
    loss_type: Literal["mse", "l1", "cosine"] = "mse",
) -> None

Initializes the GramMatrixLoss.

Parameters:

Name	Type	Description	Default
loss_type	Literal['mse', 'l1', 'cosine']	Distance metric ('mse' or 'l1'). Defaults to 'mse'.	'mse'

Raises:

Type	Description
ValueError	If an unsupported loss_type is provided.

Source code in src/shrinkai/distillation/losses/features.py

def __init__(self, loss_type: Literal["mse", "l1", "cosine"] = "mse") -> None:
    """Initializes the GramMatrixLoss.

    Args:
        loss_type: Distance metric ('mse' or 'l1'). Defaults to 'mse'.

    Raises:
        ValueError: If an unsupported `loss_type` is provided.
    """
    super().__init__()
    if loss_type not in ["mse", "l1", "cosine"]:
        raise ValueError(f"Unsupported loss_type '{loss_type}'. Use 'mse', 'l1' or 'cosine'.")

    self.loss_type = loss_type

forward

forward(
    student_outputs: Tensor | dict[str, Tensor],
    teacher_outputs: Tensor | dict[str, Tensor],
    labels: Tensor | None = None,
) -> torch.Tensor

Computes the Gram matrix distillation loss.

Parameters:

Name	Type	Description	Default
student_outputs	Tensor | dict[str, Tensor]	Tensor or dict of intermediate features [B, C, H, W].	required
teacher_outputs	Tensor | dict[str, Tensor]	Tensor or dict of intermediate features [B, C, H, W].	required
labels	Tensor | None	Ground-truth labels (ignored, kept for API compatibility).	None

Returns:

Type	Description
Tensor	torch.Tensor: Aggregated scalar loss value.

Source code in src/shrinkai/distillation/losses/features.py

def forward(
    self,
    student_outputs: torch.Tensor | dict[str, torch.Tensor],
    teacher_outputs: torch.Tensor | dict[str, torch.Tensor],
    labels: torch.Tensor | None = None,
) -> torch.Tensor:
    """Computes the Gram matrix distillation loss.

    Args:
        student_outputs: Tensor or dict of intermediate features [B, C, H, W].
        teacher_outputs: Tensor or dict of intermediate features [B, C, H, W].
        labels: Ground-truth labels (ignored, kept for API compatibility).

    Returns:
        torch.Tensor: Aggregated scalar loss value.
    """
    return self._apply_feature_distance(
        student_outputs, teacher_outputs, self._compute_distance
    )