losses

shrinkai.distillation.losses

Knowledge Distillation loss functions.

This module provides a comprehensive suite of loss functions for Knowledge Distillation (KD), structured into three main categories based on where they operate within the neural network:

• Logit-based Distillation (Response-based): Operates on the final output logits. By applying temperature scaling, these functions transfer the "dark knowledge" of the teacher (relative probabilities assigned to non-target classes) to the student, improving its generalization capabilities.

• Available Losses: HintonLoss, PureKDLoss, ReverseKLLoss, BCEKDLoss, JSDLoss.

• Feature-based Distillation: Focuses on aligning intermediate representations (hidden layers, attention maps, Gram matrices). Unlike logit-based KD, this forces the student to learn the internal reasoning process and hierarchical representations of the teacher.

• Available Losses: FeatureLoss, AttentionMapLoss, GramMatrixLoss.

• Wrappers & Composition: Orchestration modules that do not compute mathematical distances themselves. Instead, they handle tensor dimension mapping (projection) and combine multiple distillation objectives into a single criterion.

• Available Losses: ProjectedFeatureLoss, HybridLoss, CombinedLoss.

Note on Dimension Alignment: When using feature-based methods, 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. You can use the ProjectedFeatureLoss wrapper to automate this, though a custom manual projection might sometimes be preferable depending on your architecture.

Modules:

Name	Description
base	Base class for Knowledge Distillation loss functions.
features	Feature-based Knowledge Distillation loss functions.
logits	Logit-based Knowledge Distillation loss functions.
wrappers	Wrapper modules for composing and routing distillation losses.

Classes:

Name	Description
AttentionMapLoss	Distillation loss for Transformer attention maps.
BCEKDLoss	Knowledge Distillation loss for Multi-Label classification tasks.
BaseDistillationLoss	Abstract base class for all distillation loss functions.
CombinedLoss	Combines an arbitrary number of distillation losses with specific weights.
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
HintonLoss	Knowledge Distillation loss (Geoffrey Hinton et al. (2015)).
HybridLoss	Combines a primary logit-based loss and a feature-based loss using a convex combination.
JSDLoss	Jensen-Shannon Divergence (JSD) loss for Knowledge Distillation.
ProjectedFeatureLoss	Bridge between FeatureExtractors, FeatureProjectors, and Feature Losses.
PureKDLoss	Pure Knowledge Distillation loss.
ReverseKLLoss	Reverse Kullback-Leibler divergence for Knowledge Distillation (Gu et al. (2024),

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
    )

BCEKDLoss

Bases: BaseDistillationLoss

Knowledge Distillation loss for Multi-Label classification tasks.

Unlike HintonLoss which uses Softmax (classes are mutually exclusive), this loss uses Sigmoid to treat each class independently. It computes the Binary Cross Entropy (BCE) between the student's logits and the teacher's softened targets.

Equation

\[L = (1 - \alpha) \cdot \text{BCE}(z_s, y) + \alpha \cdot T^2 \cdot \text{BCE}\left(\frac{z_s}{T}, \text{sigmoid}\left(\frac{z_t}{T}\right)\right)\]

Attributes:

Name	Type	Description
temperature	float	Softening factor for logits. Must be > 0.
alpha	float	Weight balancing factor. Must be in range [0.0, 1.0].

Methods:

Name	Description
forward	Computes combined hard BCE and soft BCE losses.

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

class BCEKDLoss(BaseDistillationLoss):
    r"""Knowledge Distillation loss for Multi-Label classification tasks.

    Unlike HintonLoss which uses Softmax (classes are mutually exclusive),
    this loss uses Sigmoid to treat each class independently. It computes
    the Binary Cross Entropy (BCE) between the student's logits and the
    teacher's softened targets.

    Equation:
        $$L = (1 - \alpha) \cdot \text{BCE}(z_s, y) + \alpha \cdot T^2 \cdot \text{BCE}\left(\frac{z_s}{T}, \text{sigmoid}\left(\frac{z_t}{T}\right)\right)$$

    Attributes:
        temperature (float): Softening factor for logits. Must be > 0.
        alpha (float): Weight balancing factor. Must be in range [0.0, 1.0].
    """  # noqa: E501

    def __init__(self, temperature: float = 2.0, alpha: float = 0.5) -> None:
        super().__init__()
        if temperature <= 0:
            raise ValueError(f"Temperature must be positive, got {temperature}")
        if not (0.0 <= alpha <= 1.0):
            raise ValueError(f"Alpha must be between 0.0 and 1.0, got {alpha}")

        self.temperature = float(temperature)
        self.alpha = float(alpha)
        self.bce_with_logits = nn.BCEWithLogitsLoss()

    @expects_logits
    def forward(
        self,
        student_outputs: torch.Tensor,
        teacher_outputs: torch.Tensor,
        labels: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Computes combined hard BCE and soft BCE losses.

        Args:
            student_outputs: Raw unnormalized logits from the student model
                (Shape: [B, C] or [B, S, C]).
            teacher_outputs: Raw unnormalized logits from the teacher model
                (Shape: [B, C] or [B, S, C]).
            labels: Ground-truth labels (ignored, kept for API compatibility).

        Returns:
            torch.Tensor: Scalar loss value.

        Raises:
            ValueError: If shapes of `student_outputs` and `teacher_outputs` mismatch.
        """
        if student_outputs.shape != teacher_outputs.shape:
            raise ValueError("Shape mismatch between student and teacher logits.")

        with torch.no_grad():
            soft_targets = torch.sigmoid(teacher_outputs / self.temperature)

        kd_loss = F.binary_cross_entropy_with_logits(
            student_outputs / self.temperature, soft_targets
        )
        kd_loss = kd_loss * (self.temperature**2)

        if self.alpha == 1.0:
            return kd_loss

        if labels is None:
            raise ValueError("Ground-truth `labels` are required when `alpha` < 1.0.")

        hard_loss = self.bce_with_logits(student_outputs, labels.float())
        return (1.0 - self.alpha) * hard_loss + self.alpha * kd_loss

Methods:

forward

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

Computes combined hard BCE and soft BCE losses.

Parameters:

Name	Type	Description	Default
student_outputs	Tensor	Raw unnormalized logits from the student model (Shape: [B, C] or [B, S, C]).	required
teacher_outputs	Tensor	Raw unnormalized logits from the teacher model (Shape: [B, C] or [B, S, C]).	required
labels	Tensor | None	Ground-truth labels (ignored, kept for API compatibility).	None

Returns:

Type	Description
Tensor	torch.Tensor: Scalar loss value.

Raises:

Type	Description
ValueError	If shapes of student_outputs and teacher_outputs mismatch.

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

@expects_logits
def forward(
    self,
    student_outputs: torch.Tensor,
    teacher_outputs: torch.Tensor,
    labels: torch.Tensor | None = None,
) -> torch.Tensor:
    """Computes combined hard BCE and soft BCE losses.

    Args:
        student_outputs: Raw unnormalized logits from the student model
            (Shape: [B, C] or [B, S, C]).
        teacher_outputs: Raw unnormalized logits from the teacher model
            (Shape: [B, C] or [B, S, C]).
        labels: Ground-truth labels (ignored, kept for API compatibility).

    Returns:
        torch.Tensor: Scalar loss value.

    Raises:
        ValueError: If shapes of `student_outputs` and `teacher_outputs` mismatch.
    """
    if student_outputs.shape != teacher_outputs.shape:
        raise ValueError("Shape mismatch between student and teacher logits.")

    with torch.no_grad():
        soft_targets = torch.sigmoid(teacher_outputs / self.temperature)

    kd_loss = F.binary_cross_entropy_with_logits(
        student_outputs / self.temperature, soft_targets
    )
    kd_loss = kd_loss * (self.temperature**2)

    if self.alpha == 1.0:
        return kd_loss

    if labels is None:
        raise ValueError("Ground-truth `labels` are required when `alpha` < 1.0.")

    hard_loss = self.bce_with_logits(student_outputs, labels.float())
    return (1.0 - self.alpha) * hard_loss + self.alpha * kd_loss

BaseDistillationLoss

Bases: Module, ABC

Abstract base class for all distillation loss functions.

All concrete implementations must override the forward method.

Methods:

Name	Description
forward	Computes the distillation loss.

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

class BaseDistillationLoss(nn.Module, ABC):
    """
    Abstract base class for all distillation loss functions.

    All concrete implementations must override the `forward` method.
    """

    def __init__(self) -> None:
        super().__init__()

    @abstractmethod
    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 distillation loss.

        Args:
            student_outputs: Output tensor (or dict of activations) from the student model.
            teacher_outputs: Output tensor (or dict of activations) from the teacher model.
            labels: Ground-truth task labels (optional depending on the loss type).

        Returns:
            torch.Tensor: Scalar loss tensor for backpropagation.
        """
        pass

    def _apply_feature_distance(
        self,
        student_outputs: torch.Tensor | dict[str, torch.Tensor],
        teacher_outputs: torch.Tensor | dict[str, torch.Tensor],
        distance_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
    ) -> torch.Tensor:
        """Helper method to apply a distance function across tensors or dicts of tensors.

        Args:
            student_outputs: Tensor or dict of intermediate feature tensors.
            teacher_outputs: Tensor or dict of intermediate feature tensors.
            distance_fn: A function that computes the distance between two tensors.

        Returns:
            torch.Tensor: The averaged distance.
        """
        if isinstance(student_outputs, dict) and isinstance(teacher_outputs, dict):
            if student_outputs.keys() != teacher_outputs.keys():
                raise ValueError(
                    f"Dictionary key mismatch. Student keys: {list(student_outputs.keys())}, "
                    f"Teacher keys: {list(teacher_outputs.keys())}"
                )

            total_loss = torch.tensor(0.0, device=next(iter(student_outputs.values())).device)
            for key in student_outputs.keys():
                total_loss += distance_fn(student_outputs[key], teacher_outputs[key])

            return total_loss / len(student_outputs)

        if isinstance(student_outputs, torch.Tensor) and isinstance(teacher_outputs, torch.Tensor):
            return distance_fn(student_outputs, teacher_outputs)

        raise TypeError("Student and teacher outputs must be both Tensors or both dictionaries.")

Methods:

forward abstractmethod

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

Computes the distillation loss.

Parameters:

Name	Type	Description	Default
student_outputs	Tensor | dict[str, Tensor]	Output tensor (or dict of activations) from the student model.	required
teacher_outputs	Tensor | dict[str, Tensor]	Output tensor (or dict of activations) from the teacher model.	required
labels	Tensor | None	Ground-truth task labels (optional depending on the loss type).	None

Returns:

Type	Description
Tensor	torch.Tensor: Scalar loss tensor for backpropagation.

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

@abstractmethod
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 distillation loss.

    Args:
        student_outputs: Output tensor (or dict of activations) from the student model.
        teacher_outputs: Output tensor (or dict of activations) from the teacher model.
        labels: Ground-truth task labels (optional depending on the loss type).

    Returns:
        torch.Tensor: Scalar loss tensor for backpropagation.
    """
    pass

CombinedLoss

Bases: BaseDistillationLoss

Combines an arbitrary number of distillation losses with specific weights.

This acts as a transparent router. It passes the raw inputs (whether they are tensors or tuples) to each underlying loss.

Methods:

Name	Description
__init__	Initializes the CombinedLoss.
forward	Computes the weighted sum of all configured losses.

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

class CombinedLoss(BaseDistillationLoss):
    """Combines an arbitrary number of distillation losses with specific weights.

    This acts as a transparent router. It passes the raw inputs (whether they are
    tensors or tuples) to each underlying loss.
    """

    def __init__(self, weighted_losses: list[tuple[BaseDistillationLoss, float]]) -> None:
        """Initializes the CombinedLoss.

        Args:
            weighted_losses: A list of tuple with instantiated loss modules and their
                corresponding weights.
        """
        super().__init__()
        if not all(isinstance(item, tuple) and len(item) == 2 for item in weighted_losses):
            raise ValueError(
                "weighted_losses must be a list of tuples in the"
                "format: [(loss_module, weight), ...]"
            )
        self.losses = nn.ModuleList([loss for loss, _ in weighted_losses])
        self.weights = [weight for _, weight in weighted_losses]

    def forward(
        self,
        student_outputs: Any,
        teacher_outputs: Any,
        labels: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Computes the weighted sum of all configured losses.

        Args:
            student_outputs: Tuple of (student_logits, student_features_dict).
            teacher_outputs: Tuple of (teacher_logits, teacher_features_dict).
            labels: Ground-truth labels.

        Returns:
            torch.Tensor: Aggregated scalar loss value.
        """
        total_loss = 0.0

        for loss_fn, weight in zip(self.losses, self.weights, strict=True):
            if weight > 0.0:
                total_loss += weight * loss_fn(student_outputs, teacher_outputs, labels)

        return total_loss

Methods:

__init__

__init__(
    weighted_losses: list[
        tuple[BaseDistillationLoss, float]
    ],
) -> None

Initializes the CombinedLoss.

Parameters:

Name	Type	Description	Default
weighted_losses	list[tuple[BaseDistillationLoss, float]]	A list of tuple with instantiated loss modules and their corresponding weights.	required

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

def __init__(self, weighted_losses: list[tuple[BaseDistillationLoss, float]]) -> None:
    """Initializes the CombinedLoss.

    Args:
        weighted_losses: A list of tuple with instantiated loss modules and their
            corresponding weights.
    """
    super().__init__()
    if not all(isinstance(item, tuple) and len(item) == 2 for item in weighted_losses):
        raise ValueError(
            "weighted_losses must be a list of tuples in the"
            "format: [(loss_module, weight), ...]"
        )
    self.losses = nn.ModuleList([loss for loss, _ in weighted_losses])
    self.weights = [weight for _, weight in weighted_losses]

forward

forward(
    student_outputs: Any,
    teacher_outputs: Any,
    labels: Tensor | None = None,
) -> torch.Tensor

Computes the weighted sum of all configured losses.

Parameters:

Name	Type	Description	Default
student_outputs	Any	Tuple of (student_logits, student_features_dict).	required
teacher_outputs	Any	Tuple of (teacher_logits, teacher_features_dict).	required
labels	Tensor | None	Ground-truth labels.	None

Returns:

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

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

def forward(
    self,
    student_outputs: Any,
    teacher_outputs: Any,
    labels: torch.Tensor | None = None,
) -> torch.Tensor:
    """Computes the weighted sum of all configured losses.

    Args:
        student_outputs: Tuple of (student_logits, student_features_dict).
        teacher_outputs: Tuple of (teacher_logits, teacher_features_dict).
        labels: Ground-truth labels.

    Returns:
        torch.Tensor: Aggregated scalar loss value.
    """
    total_loss = 0.0

    for loss_fn, weight in zip(self.losses, self.weights, strict=True):
        if weight > 0.0:
            total_loss += weight * loss_fn(student_outputs, teacher_outputs, labels)

    return total_loss

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
    )

HintonLoss

Bases: BaseDistillationLoss

Knowledge Distillation loss (Geoffrey Hinton et al. (2015)).

Combines standard task loss (Cross-Entropy with hard ground-truth labels) and distillation loss (Kullback-Leibler divergence on softened probabilities produced by a teacher model at a given temperature).

Equation

\[L = (1 - \alpha) \cdot L_{ce}(z_s, y) + \alpha \cdot T^2 \cdot L_{kl}\left(\text{softmax}\left(\frac{z_s}{T}\right), \text{softmax}\left(\frac{z_t}{T}\right)\right)\]

Attributes:

Name	Type	Description
temperature	float	Softening factor for logits. Higher values produce smoother probability distributions over classes. Must be > 0.
alpha	float	Weight balancing factor between hard label loss and distillation loss. Must be in range [0.0, 1.0].

Methods:

Name	Description
__init__	Initializes the HintonLoss module.
forward	Computes combined Cross-Entropy and KD divergence losses.

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

class HintonLoss(BaseDistillationLoss):
    r"""Knowledge Distillation loss (Geoffrey Hinton et al. (2015)).

    Combines standard task loss (Cross-Entropy with hard ground-truth labels)
    and distillation loss (Kullback-Leibler divergence on softened probabilities
    produced by a teacher model at a given temperature).

    Equation:
        $$L = (1 - \alpha) \cdot L_{ce}(z_s, y) + \alpha \cdot T^2 \cdot L_{kl}\left(\text{softmax}\left(\frac{z_s}{T}\right), \text{softmax}\left(\frac{z_t}{T}\right)\right)$$

    Attributes:
        temperature (float): Softening factor for logits. Higher values produce
            smoother probability distributions over classes. Must be > 0.
        alpha (float): Weight balancing factor between hard label loss and
            distillation loss. Must be in range [0.0, 1.0].
    """  # noqa: E501

    def __init__(self, temperature: float = 4.0, alpha: float = 0.5) -> None:
        """Initializes the HintonLoss module.

        Args:
            temperature: Temperature scaling factor (T > 0). Defaults to 4.0.
            alpha: Weight for distillation loss (0.0 <= alpha <= 1.0). Defaults to 0.5.

        Raises:
            ValueError: If `temperature` <= 0 or `alpha` is not in [0.0, 1.0].
        """
        super().__init__()

        if temperature <= 0:
            raise ValueError(f"Temperature must be positive, current temperature is {temperature}")
        if not (0.0 <= alpha <= 1.0):
            raise ValueError(f"Alpha must be between 0.0 and 1.0, current alpha is {alpha}")

        self.temperature = float(temperature)
        self.alpha = float(alpha)
        self.kl_div = nn.KLDivLoss(reduction="batchmean")
        self.cross_entropy = nn.CrossEntropyLoss()

    @expects_logits
    def forward(
        self,
        student_outputs: torch.Tensor,
        teacher_outputs: torch.Tensor,
        labels: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Computes combined Cross-Entropy and KD divergence losses.

        Args:
            student_outputs: Raw unnormalized logits from the student model (Shape: [B, C]).
            teacher_outputs: Raw unnormalized logits from the teacher model (Shape: [B, C]).
            labels: Ground-truth class indices (Shape: [B]). Optional if alpha == 1.0.

        Returns:
            torch.Tensor: Weighted scalar loss value.

        Raises:
            ValueError: If `labels` is None when `alpha` < 1.0.
            ValueError: If shapes of `student_outputs` and `teacher_outputs` mismatch.
        """
        if student_outputs.shape != teacher_outputs.shape:
            raise ValueError(
                f"Shape mismatch: student logits {student_outputs.shape} "
                f"vs teacher logits {teacher_outputs.shape}"
            )

        soft_student = F.log_softmax(student_outputs / self.temperature, dim=-1)
        soft_teacher = F.softmax(teacher_outputs / self.temperature, dim=-1)
        kd_loss = self.kl_div(soft_student, soft_teacher) * (self.temperature**2)

        if self.alpha == 1.0:
            return kd_loss

        if labels is None:
            raise ValueError("Ground-truth `labels` are required when `alpha` < 1.0.")

        ce_loss = self.cross_entropy(student_outputs, labels)
        return (1.0 - self.alpha) * ce_loss + self.alpha * kd_loss

Methods:

__init__

__init__(
    temperature: float = 4.0, alpha: float = 0.5
) -> None

Initializes the HintonLoss module.

Parameters:

Name	Type	Description	Default
temperature	float	Temperature scaling factor (T > 0). Defaults to 4.0.	4.0
alpha	float	Weight for distillation loss (0.0 <= alpha <= 1.0). Defaults to 0.5.	0.5

Raises:

Type	Description
ValueError	If temperature <= 0 or alpha is not in [0.0, 1.0].

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

def __init__(self, temperature: float = 4.0, alpha: float = 0.5) -> None:
    """Initializes the HintonLoss module.

    Args:
        temperature: Temperature scaling factor (T > 0). Defaults to 4.0.
        alpha: Weight for distillation loss (0.0 <= alpha <= 1.0). Defaults to 0.5.

    Raises:
        ValueError: If `temperature` <= 0 or `alpha` is not in [0.0, 1.0].
    """
    super().__init__()

    if temperature <= 0:
        raise ValueError(f"Temperature must be positive, current temperature is {temperature}")
    if not (0.0 <= alpha <= 1.0):
        raise ValueError(f"Alpha must be between 0.0 and 1.0, current alpha is {alpha}")

    self.temperature = float(temperature)
    self.alpha = float(alpha)
    self.kl_div = nn.KLDivLoss(reduction="batchmean")
    self.cross_entropy = nn.CrossEntropyLoss()

forward

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

Computes combined Cross-Entropy and KD divergence losses.

Parameters:

Name	Type	Description	Default
student_outputs	Tensor	Raw unnormalized logits from the student model (Shape: [B, C]).	required
teacher_outputs	Tensor	Raw unnormalized logits from the teacher model (Shape: [B, C]).	required
labels	Tensor | None	Ground-truth class indices (Shape: [B]). Optional if alpha == 1.0.	None

Returns:

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

Raises:

Type	Description
ValueError	If labels is None when alpha < 1.0.
ValueError	If shapes of student_outputs and teacher_outputs mismatch.

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

@expects_logits
def forward(
    self,
    student_outputs: torch.Tensor,
    teacher_outputs: torch.Tensor,
    labels: torch.Tensor | None = None,
) -> torch.Tensor:
    """Computes combined Cross-Entropy and KD divergence losses.

    Args:
        student_outputs: Raw unnormalized logits from the student model (Shape: [B, C]).
        teacher_outputs: Raw unnormalized logits from the teacher model (Shape: [B, C]).
        labels: Ground-truth class indices (Shape: [B]). Optional if alpha == 1.0.

    Returns:
        torch.Tensor: Weighted scalar loss value.

    Raises:
        ValueError: If `labels` is None when `alpha` < 1.0.
        ValueError: If shapes of `student_outputs` and `teacher_outputs` mismatch.
    """
    if student_outputs.shape != teacher_outputs.shape:
        raise ValueError(
            f"Shape mismatch: student logits {student_outputs.shape} "
            f"vs teacher logits {teacher_outputs.shape}"
        )

    soft_student = F.log_softmax(student_outputs / self.temperature, dim=-1)
    soft_teacher = F.softmax(teacher_outputs / self.temperature, dim=-1)
    kd_loss = self.kl_div(soft_student, soft_teacher) * (self.temperature**2)

    if self.alpha == 1.0:
        return kd_loss

    if labels is None:
        raise ValueError("Ground-truth `labels` are required when `alpha` < 1.0.")

    ce_loss = self.cross_entropy(student_outputs, labels)
    return (1.0 - self.alpha) * ce_loss + self.alpha * kd_loss

HybridLoss

Bases: CombinedLoss

Combines a primary logit-based loss and a feature-based loss using a convex combination.

Equation

\[ \begin{aligned} \text{Additive } (\texttt{convex_weighting=False})\colon \quad & L = L_{\text{primary}} + \lambda \cdot L_{\text{feature}} \\\\ \text{Convex } (\texttt{convex_weighting=True})\colon \quad & L = (1 - \lambda) \cdot L_{\text{primary}} + \lambda \cdot L_{\text{feature}} \end{aligned} \]

where \(\lambda = \text{feature_weight}\).

This loss expects the models (or the FeatureExtractor wrapper) to return a tuple containing (logits, features_dict).

Methods:

Name	Description
__init__	Initializes the HybridLoss.

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

class HybridLoss(CombinedLoss):
    r"""Combines a primary logit-based loss and a feature-based loss using a convex combination.

    Equation:
        $$
        \begin{aligned}
        \text{Additive } (\texttt{convex_weighting=False})\colon \quad & L = L_{\text{primary}} + \lambda \cdot L_{\text{feature}} \\\\
        \text{Convex } (\texttt{convex_weighting=True})\colon \quad & L = (1 - \lambda) \cdot L_{\text{primary}} + \lambda \cdot L_{\text{feature}}
        \end{aligned}
        $$

        where $\lambda = \text{feature_weight}$.

    This loss expects the models (or the FeatureExtractor wrapper) to return
    a tuple containing `(logits, features_dict)`.
    """  # noqa: E501

    def __init__(
        self,
        logit_loss: BaseDistillationLoss,
        feature_loss: BaseDistillationLoss,
        feature_weight: float = 1.0,
        convex_weighting: bool = False,
    ) -> None:
        """Initializes the HybridLoss.

        Args:
            logit_loss: Loss applied to the final logits (e.g., a HintonLoss object).
            feature_loss: Loss applied to the intermediate features (e.g., a FeatureLoss object).
            feature_weight: Balancing factor (0.0 <= weight <= 1.0). Defaults to 1.
            convex_weighting: If True, applies (1-w) to primary loss and (w) to feature loss.
                If False, strictly adds (w * feature_loss) to primary loss. Defaults to False.

        Raises:
            ValueError: If `feature_weight` is not in the [0.0, 1.0] range.
        """
        if convex_weighting and not (0.0 <= feature_weight <= 1.0):
            raise ValueError(
                "For convex weighting, feature_weight must be between 0.0"
                f"and 1.0, got {feature_weight}"
            )

        if convex_weighting:
            super().__init__(
                weighted_losses=[
                    (logit_loss, 1 - feature_weight),
                    (feature_loss, feature_weight),
                ]
            )

        else:
            super().__init__(weighted_losses=[(logit_loss, 1), (feature_loss, feature_weight)])

Methods:

__init__

__init__(
    logit_loss: BaseDistillationLoss,
    feature_loss: BaseDistillationLoss,
    feature_weight: float = 1.0,
    convex_weighting: bool = False,
) -> None

Initializes the HybridLoss.

Parameters:

Name	Type	Description	Default
logit_loss	BaseDistillationLoss	Loss applied to the final logits (e.g., a HintonLoss object).	required
feature_loss	BaseDistillationLoss	Loss applied to the intermediate features (e.g., a FeatureLoss object).	required
feature_weight	float	Balancing factor (0.0 <= weight <= 1.0). Defaults to 1.	1.0
convex_weighting	bool	If True, applies (1-w) to primary loss and (w) to feature loss. If False, strictly adds (w * feature_loss) to primary loss. Defaults to False.	False

Raises:

Type	Description
ValueError	If feature_weight is not in the [0.0, 1.0] range.

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

def __init__(
    self,
    logit_loss: BaseDistillationLoss,
    feature_loss: BaseDistillationLoss,
    feature_weight: float = 1.0,
    convex_weighting: bool = False,
) -> None:
    """Initializes the HybridLoss.

    Args:
        logit_loss: Loss applied to the final logits (e.g., a HintonLoss object).
        feature_loss: Loss applied to the intermediate features (e.g., a FeatureLoss object).
        feature_weight: Balancing factor (0.0 <= weight <= 1.0). Defaults to 1.
        convex_weighting: If True, applies (1-w) to primary loss and (w) to feature loss.
            If False, strictly adds (w * feature_loss) to primary loss. Defaults to False.

    Raises:
        ValueError: If `feature_weight` is not in the [0.0, 1.0] range.
    """
    if convex_weighting and not (0.0 <= feature_weight <= 1.0):
        raise ValueError(
            "For convex weighting, feature_weight must be between 0.0"
            f"and 1.0, got {feature_weight}"
        )

    if convex_weighting:
        super().__init__(
            weighted_losses=[
                (logit_loss, 1 - feature_weight),
                (feature_loss, feature_weight),
            ]
        )

    else:
        super().__init__(weighted_losses=[(logit_loss, 1), (feature_loss, feature_weight)])

JSDLoss

Bases: BaseDistillationLoss

Jensen-Shannon Divergence (JSD) loss for Knowledge Distillation.

JSD is a symmetric and bounded alternative to the standard KL Divergence. It computes the divergence of both distributions from their average distribution M. This boundedness prevents gradient explosion, especially early in training when the student's predictions might diverge heavily from the teacher's.

Equation

\[ \begin{aligned} M &= 0.5 \cdot (P_{student} + P_{teacher}) \\ L_{JSD} &= 0.5 \cdot \text{KL}(P_{student} \parallel M) + 0.5 \cdot \text{KL}(P_{teacher} \parallel M) \end{aligned} \]

Attributes:

Name	Type	Description
temperature	float	Softening factor for logits. Must be > 0.

Methods:

Name	Description
__init__	Initializes the JSDLoss module.
forward	Computes the JSD between softened student and teacher distributions.

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

class JSDLoss(BaseDistillationLoss):
    r"""Jensen-Shannon Divergence (JSD) loss for Knowledge Distillation.

    JSD is a symmetric and bounded alternative to the standard KL Divergence.
    It computes the divergence of both distributions from their average distribution M.
    This boundedness prevents gradient explosion, especially early in training when
    the student's predictions might diverge heavily from the teacher's.

    Equation:
        $$
        \begin{aligned}
        M &= 0.5 \cdot (P_{student} + P_{teacher}) \\
        L_{JSD} &= 0.5 \cdot \text{KL}(P_{student} \parallel M) + 0.5 \cdot \text{KL}(P_{teacher} \parallel M)
        \end{aligned}
        $$

    Attributes:
        temperature (float): Softening factor for logits. Must be > 0.
    """  # noqa: E501

    def __init__(self, temperature: float = 4.0) -> None:
        """Initializes the JSDLoss module.

        Args:
            temperature: Softening factor for logits (T > 0). Defaults to 4.0.

        Raises:
            ValueError: If `temperature` <= 0.
        """
        super().__init__()
        if temperature <= 0:
            raise ValueError(f"Temperature must be positive, got {temperature}")

        self.temperature = float(temperature)
        self.kl_div = nn.KLDivLoss(reduction="batchmean")

    @expects_logits
    def forward(
        self,
        student_outputs: torch.Tensor,
        teacher_outputs: torch.Tensor,
        labels: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Computes the JSD between softened student and teacher distributions.

         Args:
            student_outputs: Raw unnormalized logits from the student model (Shape: [B, C] or [B, S, C]).
            teacher_outputs: Raw unnormalized logits from the teacher model (Shape: [B, C] or [B, S, C]).
            labels: Ground-truth labels (ignored, kept for API compatibility).

        Returns:
            torch.Tensor: Scalar loss value.

        Raises:
            ValueError: If shapes of `student_outputs` and `teacher_outputs` mismatch.
        """  # noqa: E501
        if student_outputs.shape != teacher_outputs.shape:
            raise ValueError("Shape mismatch between student and teacher logits.")

        p_s = F.softmax(student_outputs / self.temperature, dim=-1)
        p_t = F.softmax(teacher_outputs / self.temperature, dim=-1)

        m = 0.5 * (p_s + p_t)

        log_m = torch.log(m.clamp(min=1e-8))
        kl_s_m = self.kl_div(input=log_m, target=p_s)
        kl_t_m = self.kl_div(input=log_m, target=p_t)

        jsd = 0.5 * (kl_s_m + kl_t_m)

        return jsd * (self.temperature**2)

Methods:

__init__

__init__(temperature: float = 4.0) -> None

Initializes the JSDLoss module.

Parameters:

Name	Type	Description	Default
temperature	float	Softening factor for logits (T > 0). Defaults to 4.0.	4.0

Raises:

Type	Description
ValueError	If temperature <= 0.

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

def __init__(self, temperature: float = 4.0) -> None:
    """Initializes the JSDLoss module.

    Args:
        temperature: Softening factor for logits (T > 0). Defaults to 4.0.

    Raises:
        ValueError: If `temperature` <= 0.
    """
    super().__init__()
    if temperature <= 0:
        raise ValueError(f"Temperature must be positive, got {temperature}")

    self.temperature = float(temperature)
    self.kl_div = nn.KLDivLoss(reduction="batchmean")

forward

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

Computes the JSD between softened student and teacher distributions.

Args: student_outputs: Raw unnormalized logits from the student model (Shape: [B, C] or [B, S, C]). teacher_outputs: Raw unnormalized logits from the teacher model (Shape: [B, C] or [B, S, C]). labels: Ground-truth labels (ignored, kept for API compatibility).

Returns:

Type	Description
Tensor	torch.Tensor: Scalar loss value.

Raises:

Type	Description
ValueError	If shapes of student_outputs and teacher_outputs mismatch.

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

@expects_logits
def forward(
    self,
    student_outputs: torch.Tensor,
    teacher_outputs: torch.Tensor,
    labels: torch.Tensor | None = None,
) -> torch.Tensor:
    """Computes the JSD between softened student and teacher distributions.

     Args:
        student_outputs: Raw unnormalized logits from the student model (Shape: [B, C] or [B, S, C]).
        teacher_outputs: Raw unnormalized logits from the teacher model (Shape: [B, C] or [B, S, C]).
        labels: Ground-truth labels (ignored, kept for API compatibility).

    Returns:
        torch.Tensor: Scalar loss value.

    Raises:
        ValueError: If shapes of `student_outputs` and `teacher_outputs` mismatch.
    """  # noqa: E501
    if student_outputs.shape != teacher_outputs.shape:
        raise ValueError("Shape mismatch between student and teacher logits.")

    p_s = F.softmax(student_outputs / self.temperature, dim=-1)
    p_t = F.softmax(teacher_outputs / self.temperature, dim=-1)

    m = 0.5 * (p_s + p_t)

    log_m = torch.log(m.clamp(min=1e-8))
    kl_s_m = self.kl_div(input=log_m, target=p_s)
    kl_t_m = self.kl_div(input=log_m, target=p_t)

    jsd = 0.5 * (kl_s_m + kl_t_m)

    return jsd * (self.temperature**2)

ProjectedFeatureLoss

Bases: BaseDistillationLoss

Bridge between FeatureExtractors, FeatureProjectors, and Feature Losses.

When using FeatureExtractor, the model outputs a tuple: (logits, features_dict). However, feature losses (like FeatureLoss, AttentionMapLoss) expect pure dictionaries or tensors.

This wrapper seamlessly unpacks the tuples, applies the FeatureProjector to align the student's feature dimensions with the teacher's, and computes the underlying feature loss.

Attributes:

Name	Type	Description
projector	Module	The module responsible for projecting student features.
feature_loss	BaseDistillationLoss	The loss function to apply to the aligned features.
project_teacher	bool	If False (default), the projector is applied to the student's features to match the teacher's larger dimensions. If True, the projector is applied to the teacher's features to downscale them (e.g., filtering a 12-head teacher down to match a 2-head student in Transformer attention distillation).

Methods:

Name	Description
__init__	Initializes the ProjectedFeatureLoss.
forward	Unpacks outputs, projects features (student or teacher), and computes the loss.

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

class ProjectedFeatureLoss(BaseDistillationLoss):
    """Bridge between FeatureExtractors, FeatureProjectors, and Feature Losses.

    When using `FeatureExtractor`, the model outputs a tuple: `(logits, features_dict)`.
    However, feature losses (like `FeatureLoss`, `AttentionMapLoss`) expect pure
    dictionaries or tensors.

    This wrapper seamlessly unpacks the tuples, applies the `FeatureProjector` to
    align the student's feature dimensions with the teacher's, and computes the
    underlying feature loss.

    Attributes:
        projector (nn.Module): The module responsible for projecting student features.
        feature_loss (BaseDistillationLoss): The loss function to apply to the aligned features.
        project_teacher (bool): If False (default), the projector is applied to the
            student's features to match the teacher's larger dimensions.
            If True, the projector is applied to the teacher's features to
            downscale them (e.g., filtering a 12-head teacher down to match
            a 2-head student in Transformer attention distillation).
    """

    def __init__(
        self,
        projector: nn.Module,
        feature_loss: BaseDistillationLoss,
        project_teacher: bool = False,
    ) -> None:
        """Initializes the ProjectedFeatureLoss.

        Args:
            projector: An instantiated FeatureProjector.
            feature_loss: An instantiated feature distillation loss (e.g., FeatureLoss).
            project_teacher: Either the projector is applied to the teacher or not.
        """
        super().__init__()
        self.projector = projector
        self.feature_loss = feature_loss
        self.project_teacher = project_teacher

    def forward(
        self,
        student_outputs: tuple[torch.Tensor, dict[str, torch.Tensor]],
        teacher_outputs: tuple[torch.Tensor, dict[str, torch.Tensor]],
        labels: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Unpacks outputs, projects features (student or teacher), and computes the loss."""
        if not isinstance(student_outputs, tuple) or not isinstance(teacher_outputs, tuple):
            raise ValueError(
                "ProjectedFeatureLoss expects outputs to be a tuple of (logits, features_dict). "
                "Ensure your models are wrapped with `FeatureExtractor`."
            )

        _, s_features_dict = student_outputs
        _, t_features_dict = teacher_outputs

        # 1. projector aligns dimensions
        if self.project_teacher:
            s_current = dict(s_features_dict)
            t_current = self.projector(t_features_dict)
        else:
            s_current = self.projector(s_features_dict)
            t_current = dict(t_features_dict)

        # 2. Spatial / temporal alignment
        for key in s_current.keys():
            if key not in t_current:
                continue

            s_tensor = s_current[key]
            t_tensor = t_current[key]

            # A : vision (CNN) -> 4D tensors [Batch, Channels, Height, Width]
            if s_tensor.dim() == 4 and t_tensor.dim() == 4:
                if s_tensor.shape[2:] != t_tensor.shape[2:]:
                    s_current[key] = F.adaptive_avg_pool2d(s_tensor, output_size=t_tensor.shape[2:])

            # B : NLP / series (Transformers, RNN) -> 3D tensors [Batch, Seq_Len, Hidden_Dim]
            elif s_tensor.dim() == 3 and t_tensor.dim() == 3:
                # In PyTorch NLP, we should have [B, L, D], but if L is different
                if s_tensor.shape[1] != t_tensor.shape[1]:
                    # adaptive_avg_pool1d expects [Batch, Channels, Length]
                    # transpose [B, L, D] -> [B, D, L], then pool, then re-transpose
                    s_transposed = s_tensor.transpose(1, 2)
                    target_length = t_tensor.shape[1]
                    s_pooled = F.adaptive_avg_pool1d(s_transposed, output_size=target_length)
                    s_current[key] = s_pooled.transpose(1, 2)

        # 3. final loss (MSE, KL, etc.)
        return self.feature_loss(
            student_outputs=s_current, teacher_outputs=t_current, labels=labels
        )

Methods:

__init__

__init__(
    projector: Module,
    feature_loss: BaseDistillationLoss,
    project_teacher: bool = False,
) -> None

Initializes the ProjectedFeatureLoss.

Parameters:

Name	Type	Description	Default
projector	Module	An instantiated FeatureProjector.	required
feature_loss	BaseDistillationLoss	An instantiated feature distillation loss (e.g., FeatureLoss).	required
project_teacher	bool	Either the projector is applied to the teacher or not.	False

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

def __init__(
    self,
    projector: nn.Module,
    feature_loss: BaseDistillationLoss,
    project_teacher: bool = False,
) -> None:
    """Initializes the ProjectedFeatureLoss.

    Args:
        projector: An instantiated FeatureProjector.
        feature_loss: An instantiated feature distillation loss (e.g., FeatureLoss).
        project_teacher: Either the projector is applied to the teacher or not.
    """
    super().__init__()
    self.projector = projector
    self.feature_loss = feature_loss
    self.project_teacher = project_teacher

forward

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

Unpacks outputs, projects features (student or teacher), and computes the loss.

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

def forward(
    self,
    student_outputs: tuple[torch.Tensor, dict[str, torch.Tensor]],
    teacher_outputs: tuple[torch.Tensor, dict[str, torch.Tensor]],
    labels: torch.Tensor | None = None,
) -> torch.Tensor:
    """Unpacks outputs, projects features (student or teacher), and computes the loss."""
    if not isinstance(student_outputs, tuple) or not isinstance(teacher_outputs, tuple):
        raise ValueError(
            "ProjectedFeatureLoss expects outputs to be a tuple of (logits, features_dict). "
            "Ensure your models are wrapped with `FeatureExtractor`."
        )

    _, s_features_dict = student_outputs
    _, t_features_dict = teacher_outputs

    # 1. projector aligns dimensions
    if self.project_teacher:
        s_current = dict(s_features_dict)
        t_current = self.projector(t_features_dict)
    else:
        s_current = self.projector(s_features_dict)
        t_current = dict(t_features_dict)

    # 2. Spatial / temporal alignment
    for key in s_current.keys():
        if key not in t_current:
            continue

        s_tensor = s_current[key]
        t_tensor = t_current[key]

        # A : vision (CNN) -> 4D tensors [Batch, Channels, Height, Width]
        if s_tensor.dim() == 4 and t_tensor.dim() == 4:
            if s_tensor.shape[2:] != t_tensor.shape[2:]:
                s_current[key] = F.adaptive_avg_pool2d(s_tensor, output_size=t_tensor.shape[2:])

        # B : NLP / series (Transformers, RNN) -> 3D tensors [Batch, Seq_Len, Hidden_Dim]
        elif s_tensor.dim() == 3 and t_tensor.dim() == 3:
            # In PyTorch NLP, we should have [B, L, D], but if L is different
            if s_tensor.shape[1] != t_tensor.shape[1]:
                # adaptive_avg_pool1d expects [Batch, Channels, Length]
                # transpose [B, L, D] -> [B, D, L], then pool, then re-transpose
                s_transposed = s_tensor.transpose(1, 2)
                target_length = t_tensor.shape[1]
                s_pooled = F.adaptive_avg_pool1d(s_transposed, output_size=target_length)
                s_current[key] = s_pooled.transpose(1, 2)

    # 3. final loss (MSE, KL, etc.)
    return self.feature_loss(
        student_outputs=s_current, teacher_outputs=t_current, labels=labels
    )

PureKDLoss

Bases: HintonLoss

Pure Knowledge Distillation loss.

Implement distillation loss (Kullback-Leibler divergence on softened probabilities produced by a teacher model at a given temperature). In fact, a PureKDLossobject is just a HintonLoss object with alpha set to 1.

Equation

\[L = T^2 \cdot L_{kl}\left(\text{softmax}\left(\frac{z_s}{T}\right), \text{softmax}\left(\frac{z_t}{T}\right)\right)\]

Attributes:

Name	Type	Description
temperature	float	Softening factor for logits. Higher values produce smoother probability distributions over classes. Must be > 0.

Methods:

Name	Description
__init__	Initializes the PureKDLoss module.

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

class PureKDLoss(HintonLoss):
    r"""Pure Knowledge Distillation loss.

    Implement distillation loss (Kullback-Leibler divergence on softened probabilities
    produced by a teacher model at a given temperature). In fact, a `PureKDLoss`object
    is just a `HintonLoss` object with `alpha` set to 1.

    Equation:
        $$L = T^2 \cdot L_{kl}\left(\text{softmax}\left(\frac{z_s}{T}\right), \text{softmax}\left(\frac{z_t}{T}\right)\right)$$

    Attributes:
        temperature (float): Softening factor for logits. Higher values produce
            smoother probability distributions over classes. Must be > 0.
    """  # noqa: E501

    def __init__(self, temperature: float = 4.0) -> None:
        """Initializes the PureKDLoss module.

        Args:
            temperature: Temperature scaling factor (T > 0). Defaults to 4.0.

        Raises:
            ValueError: If `temperature` <= 0.
        """
        super().__init__(temperature=temperature, alpha=1.0)

Methods:

__init__

__init__(temperature: float = 4.0) -> None

Initializes the PureKDLoss module.

Parameters:

Name	Type	Description	Default
temperature	float	Temperature scaling factor (T > 0). Defaults to 4.0.	4.0

Raises:

Type	Description
ValueError	If temperature <= 0.

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

def __init__(self, temperature: float = 4.0) -> None:
    """Initializes the PureKDLoss module.

    Args:
        temperature: Temperature scaling factor (T > 0). Defaults to 4.0.

    Raises:
        ValueError: If `temperature` <= 0.
    """
    super().__init__(temperature=temperature, alpha=1.0)

ReverseKLLoss

Bases: BaseDistillationLoss

Reverse Kullback-Leibler divergence for Knowledge Distillation (Gu et al. (2024), MiniLLM).

Standard KD (Forward KL) computes KL(P_teacher || P_student), which is "mode-covering". Reverse KL computes KL(P_student || P_teacher), which is "mode-seeking".

For Large Language Models (LLMs), Reverse KL is highly effective because it strongly penalizes the student for assigning high probabilities to tokens that the teacher considers unlikely, thereby reducing hallucinations and degeneration.

Equation

\[L = T^2 \cdot \sum \left( P_s \cdot (\log(P_s) - \log(P_t)) \right)\]

Attributes:

Name	Type	Description
temperature	float	Softening factor for logits. Must be > 0. Note: In LLM distillation, temperature is often set close to 1.0.

Methods:

Name	Description
__init__	Initializes the ReverseKLLoss module.
forward	Computes the Reverse KL divergence loss.

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

class ReverseKLLoss(BaseDistillationLoss):
    r"""Reverse Kullback-Leibler divergence for Knowledge Distillation (Gu et al. (2024),
    MiniLLM).

    Standard KD (Forward KL) computes KL(P_teacher || P_student), which is "mode-covering".
    Reverse KL computes KL(P_student || P_teacher), which is "mode-seeking".

    For Large Language Models (LLMs), Reverse KL is highly effective because it strongly
    penalizes the student for assigning high probabilities to tokens that the teacher
    considers unlikely, thereby reducing hallucinations and degeneration.

    Equation:
        $$L = T^2 \cdot \sum \left( P_s \cdot (\log(P_s) - \log(P_t)) \right)$$

    Attributes:
        temperature (float): Softening factor for logits. Must be > 0.
            Note: In LLM distillation, temperature is often set close to 1.0.
    """  # noqa: E501

    def __init__(self, temperature: float = 1.0) -> None:
        """Initializes the ReverseKLLoss module.

        Args:
            temperature: Temperature scaling factor (T > 0). Defaults to 1.0.

        Raises:
            ValueError: If `temperature` <= 0.
        """
        super().__init__()

        if temperature <= 0:
            raise ValueError(f"Temperature must be positive, current temperature is {temperature}")

        self.temperature = float(temperature)
        self.kl_div = nn.KLDivLoss(reduction="batchmean", log_target=True)

    @expects_logits
    def forward(
        self,
        student_outputs: torch.Tensor,
        teacher_outputs: torch.Tensor,
        labels: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Computes the Reverse KL divergence loss.

        Args:
            student_outputs: Raw unnormalized logits from the student model
                (Shape: [B, C] or [B, S, C]).
            teacher_outputs: Raw unnormalized logits from the teacher model
                (Shape: [B, C] or [B, S, C]).
            labels: Ground-truth labels (ignored, kept for API compatibility).

        Returns:
            torch.Tensor: Scalar loss value.

        Raises:
            ValueError: If shapes of `student_outputs` and `teacher_outputs` mismatch.
        """
        if student_outputs.shape != teacher_outputs.shape:
            raise ValueError(
                f"Shape mismatch: student logits {student_outputs.shape} "
                f"vs teacher logits {teacher_outputs.shape}"
            )

        # IMPORTANT: Do not swap `input` and `target`!
        # PyTorch's KLDivLoss computes KL(target || input).
        # Standard KD (Forward KL) passes `target=teacher, input=student` -> KL(teacher || student).
        # For Reverse KL, we want KL(student || teacher), so we MUST pass:
        # `target=student, input=teacher`.
        log_soft_student = F.log_softmax(student_outputs / self.temperature, dim=-1)
        log_soft_teacher = F.log_softmax(teacher_outputs / self.temperature, dim=-1)
        kd_loss = self.kl_div(input=log_soft_teacher, target=log_soft_student)
        return kd_loss * (self.temperature**2)

Methods:

__init__

__init__(temperature: float = 1.0) -> None

Initializes the ReverseKLLoss module.

Parameters:

Name	Type	Description	Default
temperature	float	Temperature scaling factor (T > 0). Defaults to 1.0.	1.0

Raises:

Type	Description
ValueError	If temperature <= 0.

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

def __init__(self, temperature: float = 1.0) -> None:
    """Initializes the ReverseKLLoss module.

    Args:
        temperature: Temperature scaling factor (T > 0). Defaults to 1.0.

    Raises:
        ValueError: If `temperature` <= 0.
    """
    super().__init__()

    if temperature <= 0:
        raise ValueError(f"Temperature must be positive, current temperature is {temperature}")

    self.temperature = float(temperature)
    self.kl_div = nn.KLDivLoss(reduction="batchmean", log_target=True)

forward

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

Computes the Reverse KL divergence loss.

Parameters:

Name	Type	Description	Default
student_outputs	Tensor	Raw unnormalized logits from the student model (Shape: [B, C] or [B, S, C]).	required
teacher_outputs	Tensor	Raw unnormalized logits from the teacher model (Shape: [B, C] or [B, S, C]).	required
labels	Tensor | None	Ground-truth labels (ignored, kept for API compatibility).	None

Returns:

Type	Description
Tensor	torch.Tensor: Scalar loss value.

Raises:

Type	Description
ValueError	If shapes of student_outputs and teacher_outputs mismatch.

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

@expects_logits
def forward(
    self,
    student_outputs: torch.Tensor,
    teacher_outputs: torch.Tensor,
    labels: torch.Tensor | None = None,
) -> torch.Tensor:
    """Computes the Reverse KL divergence loss.

    Args:
        student_outputs: Raw unnormalized logits from the student model
            (Shape: [B, C] or [B, S, C]).
        teacher_outputs: Raw unnormalized logits from the teacher model
            (Shape: [B, C] or [B, S, C]).
        labels: Ground-truth labels (ignored, kept for API compatibility).

    Returns:
        torch.Tensor: Scalar loss value.

    Raises:
        ValueError: If shapes of `student_outputs` and `teacher_outputs` mismatch.
    """
    if student_outputs.shape != teacher_outputs.shape:
        raise ValueError(
            f"Shape mismatch: student logits {student_outputs.shape} "
            f"vs teacher logits {teacher_outputs.shape}"
        )

    # IMPORTANT: Do not swap `input` and `target`!
    # PyTorch's KLDivLoss computes KL(target || input).
    # Standard KD (Forward KL) passes `target=teacher, input=student` -> KL(teacher || student).
    # For Reverse KL, we want KL(student || teacher), so we MUST pass:
    # `target=student, input=teacher`.
    log_soft_student = F.log_softmax(student_outputs / self.temperature, dim=-1)
    log_soft_teacher = F.log_softmax(teacher_outputs / self.temperature, dim=-1)
    kd_loss = self.kl_div(input=log_soft_teacher, target=log_soft_student)
    return kd_loss * (self.temperature**2)