Parametric Distiller

distiller

CLASS	DESCRIPTION
Distiller	Universal knowledge distillation framework for

Classes

Distiller(model: dict[str, ModelType], loss_fn: str = 'distillation', loss_kwargs: dict | None = None, optimizer_fn: type[torch.optim.Optimizer] | str | None = None, optimizer_kwargs: dict | None = None, lr_scheduler_fn: Any | str | None = None, lr_scheduler_kwargs: dict | None = None, device: Literal['cuda', 'cpu'] = 'cuda')

Bases: Parametric

Universal knowledge distillation framework for Parametric methods.

Automatically adapts to any teacher model type and uses appropriate distillation strategies to transfer knowledge to smaller student models.

PARAMETER	DESCRIPTION
model	A dictionary containing 'student' and 'teacher' keys. TYPE: dict[str, ModelType]
loss_fn	Distillation loss function. Options include: "distillation": DistillationLoss "feature_matching": FeatureMatchingLoss "spectral_preservation": SpectralPreservationLoss TYPE: str, by default "distillation" DEFAULT: 'distillation'
loss_kwargs	Additional keyword arguments for the loss function. See distillation losses for available options. TYPE: dict, optional, by default None DEFAULT: None
optimizer_fn	Optimizer specification for training. None: Uses Adam optimizer (default) str: Optimizer name from torch.optim (e.g., "AdamW", "SGD") type[torch.optim.Optimizer]: Optimizer class (supports third-party) TYPE: type[torch.optim.Optimizer] | str | None, optional, by default None DEFAULT: None
optimizer_kwargs	Keyword arguments for optimizer instantiation (e.g., {"lr": 1e-3, "weight_decay": 1e-4}). TYPE: dict | None, optional, by default None DEFAULT: None
lr_scheduler_fn	Learning rate scheduler specification. None: No scheduler (default) str: Scheduler name from torch.optim.lr_scheduler (e.g., "ReduceLROnPlateau") type: Scheduler class TYPE: Any | str | None, optional, by default None DEFAULT: None
lr_scheduler_kwargs	Keyword arguments for scheduler instantiation. TYPE: dict | None, optional, by default None DEFAULT: None
device	Device to run the distillation on. TYPE: Literal["cuda", "cpu"], optional, by default "cuda" DEFAULT: 'cuda'

Examples:

Basic distillation:

>>> import torch
>>> from spectre.core import Model
>>> from spectre.parametric import Distiller
>>> # Create teacher and student models
>>> teacher = Model(in_features=20, out_features=10, multipliers=[2.0, 2.0])
>>> student = teacher.compress(ratio=0.5, preserve_io=True)
>>>
>>> distiller = Distiller(
...     model={"teacher": teacher, "student": student},
...     loss_fn="distillation",
...     loss_kwargs={"temperature": 4.0, "alpha": 0.7},
... )

Distillation with temperature and alpha annealing:

>>> from spectre.utils.callbacks import ParameterAnnealingCallback
>>> distiller = Distiller(
...     model={"teacher": teacher, "student": student},
...     loss_fn="distillation",
...     loss_kwargs={"temperature": 4.0, "alpha": 0.7},
... )
>>> # Create separate callbacks for temperature and alpha annealing
>>> temp_callback = ParameterAnnealingCallback(
...     target_attr="temperature",
...     start_temperature=4.0,
...     end_temperature=1.0,
...     schedule="linear",
... )
>>> alpha_callback = ParameterAnnealingCallback(
...     target_attr="alpha",
...     start_temperature=0.7,
...     end_temperature=0.3,
...     schedule="linear",
... )
>>> distiller.fit(data, n_epochs=100, callbacks=[temp_callback, alpha_callback])

Spectral preservation with custom kernel:

>>> from spectre.kernel import t_kernel
>>> from spectre.pairwise_distance import pairwise_distance_euclidean
>>> distiller_spectral = Distiller(
...     model={"teacher": teacher, "student": student},
...     loss_fn="spectral_preservation",
...     loss_kwargs={
...         "eigenvalue_weight": 1.0,
...         "eigengap_weight": 0.5,
...         "distance_fn": pairwise_distance_euclidean,
...         "kernel_fn": t_kernel,
...         "kernel_kwargs": {
...             "alpha": torch.tensor(1.0),
...             "beta": torch.tensor(3.0),
...         },
...     },
... )

Notes

Parameter Annealing:

To anneal loss parameters (temperature, alpha) during training, use ParameterAnnealingCallback. This callback supports multiple schedules (linear, exponential, cosine) and can target any loss parameter. Multiple callbacks can be used simultaneously to anneal different parameters independently.

METHOD	DESCRIPTION
loss	Compute distillation loss between student and teacher outputs.
forward_step	Forward pass through the student network.
score_step	Compute distillation loss for evaluation without training status checks.
training_step	Compute distillation loss for a training batch.
compare_student_teacher	Compare teacher and student outputs on given input.

Source code in spectre/parametric/distiller.py

def __init__(
    self,
    model: dict[str, ModelType],
    loss_fn: str = "distillation",
    loss_kwargs: dict | None = None,
    optimizer_fn: type[torch.optim.Optimizer] | str | None = None,
    optimizer_kwargs: dict | None = None,
    lr_scheduler_fn: Any | str | None = None,
    lr_scheduler_kwargs: dict | None = None,
    device: Literal["cuda", "cpu"] = "cuda",
) -> None:
    if not isinstance(model, dict):
        raise TypeError(
            "Distiller's model must be a dictionary containing 'student' "
            "and 'teacher' keys."
        )
    if "student" not in model or "teacher" not in model:
        raise ValueError(
            "Distiller's model dict must contain both 'student' and 'teacher' keys."
        )

    super().__init__(
        model=model,
        device=device,
        optimizer_fn=optimizer_fn,
        optimizer_kwargs=optimizer_kwargs,
        lr_scheduler_fn=lr_scheduler_fn,
        lr_scheduler_kwargs=lr_scheduler_kwargs,
    )

    self.student = self.model_dict["student"]
    self.teacher = self.model_dict["teacher"]

    if hasattr(self.teacher, "to"):
        self.teacher = self.teacher.to(device)

    if hasattr(self.teacher, "is_trained") and not self.teacher.is_trained:
        raise ValueError("Teacher must be trained before distillation.")

    self.teacher.eval()
    for param in self.teacher.parameters():
        param.requires_grad = False

    if loss_kwargs is None or not isinstance(loss_kwargs, dict):
        loss_kwargs = {}

    if loss_fn == "distillation":
        self.loss_fn = DistillationLoss(**loss_kwargs)
    elif loss_fn == "feature_matching":
        self.loss_fn = FeatureMatchingLoss(**loss_kwargs)
    elif loss_fn == "spectral_preservation":
        self.loss_fn = SpectralPreservationLoss(**loss_kwargs)
    else:
        raise ValueError(
            f"Invalid distillation loss '{loss_fn}'. Must be one of "
            "['distillation', 'feature_matching', 'spectral_preservation']."
        )

Functions

loss(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Compute distillation loss between student and teacher outputs.

Performs forward pass through both student and teacher networks, then computes the selected distillation loss ('distillation', 'feature_matching', or 'spectral_preservation').

PARAMETER	DESCRIPTION
X	Input data tensor to process through both networks. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Sample weights for weighted loss computation. TYPE: torch.Tensor of shape (n_samples,) or None, by default None DEFAULT: None
target	Optional ground truth targets for supervised distillation. Used in combination with teacher knowledge when available. TYPE: torch.Tensor or None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Scalar distillation loss value for optimization.

Source code in spectre/parametric/distiller.py

def loss(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute distillation loss between student and teacher outputs.

    Performs forward pass through both student and teacher networks, then
    computes the selected distillation loss ('distillation', 'feature_matching',
    or 'spectral_preservation').

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data tensor to process through both networks.

    weights : torch.Tensor of shape (n_samples,) or None, by default None
        Sample weights for weighted loss computation.

    target : torch.Tensor or None, by default None
        Optional ground truth targets for supervised distillation. Used in
        combination with teacher knowledge when available.

    Returns
    -------
    torch.Tensor
        Scalar distillation loss value for optimization.
    """
    s_batch = self.forward_step(X)

    with torch.no_grad():
        t_batch = self.teacher(X)

    # Create context dictionary for loss computation
    context = {"S": s_batch, "T": t_batch}
    if weights is not None:
        context["weights"] = weights
    if target is not None:
        context["target"] = target

    loss = self.loss_fn(context)

    return loss

forward_step(X: torch.Tensor) -> torch.Tensor

Forward pass through the student network.

Processes input data through the student model to generate predictions or representations for distillation training.

PARAMETER	DESCRIPTION
X	Input data to process through the student network. TYPE: torch.Tensor of shape (n_samples, in_features)

RETURNS	DESCRIPTION
torch.Tensor of shape (n_samples, out_features)	Student model output.

Source code in spectre/parametric/distiller.py

def forward_step(self, X: torch.Tensor) -> torch.Tensor:
    """
    Forward pass through the student network.

    Processes input data through the student model to generate predictions
    or representations for distillation training.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data to process through the student network.

    Returns
    -------
    torch.Tensor of shape (n_samples, out_features)
        Student model output.
    """
    X = self.student(X)
    return X

score_step(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Compute distillation loss for evaluation without training status checks.

This method provides direct access to the distillation loss computation for evaluation purposes, bypassing the is_trained requirement of the inherited score method.

PARAMETER	DESCRIPTION
X	Input data to process through student and teacher networks. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Sample weights for weighted loss computation. If None, uniform weighting is applied. TYPE: torch.Tensor of shape (n_samples,) or None, by default None DEFAULT: None
target	Ground truth targets for supervised distillation evaluation. TYPE: torch.Tensor or None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Scalar distillation loss value.

Source code in spectre/parametric/distiller.py

def score_step(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute distillation loss for evaluation without training status checks.

    This method provides direct access to the distillation loss computation
    for evaluation purposes, bypassing the `is_trained` requirement
    of the inherited `score` method.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data to process through student and teacher networks.

    weights : torch.Tensor of shape (n_samples,) or None, by default None
        Sample weights for weighted loss computation. If None, uniform
        weighting is applied.

    target : torch.Tensor or None, by default None
        Ground truth targets for supervised distillation evaluation.

    Returns
    -------
    torch.Tensor
        Scalar distillation loss value.
    """
    return self.loss(X=X, weights=weights, target=target)

training_step(batch: DataBatch, batch_idx: int) -> torch.Tensor

Compute distillation loss for a training batch.

PARAMETER	DESCRIPTION
batch	Batch of data containing 'data', 'weights', and optionally 'target' attributes. TYPE: DataBatch
batch_idx	Batch index. TYPE: int

RETURNS	DESCRIPTION
Tensor	Scalar distillation loss value for backpropagation.

Source code in spectre/parametric/distiller.py

def training_step(self, batch: DataBatch, batch_idx: int) -> torch.Tensor:
    """
    Compute distillation loss for a training batch.

    Parameters
    ----------
    batch : DataBatch
        Batch of data containing 'data', 'weights', and optionally 'target'
        attributes.

    batch_idx : int
        Batch index.

    Returns
    -------
    torch.Tensor
        Scalar distillation loss value for backpropagation.
    """
    loss = self.loss(X=batch.data, weights=batch.weights, target=batch.target)

    self.log_training_metrics(loss)

    return loss

compare_student_teacher(X: torch.Tensor) -> dict[str, torch.Tensor]

Compare teacher and student outputs on given input.

PARAMETER	DESCRIPTION
x	Input data for comparison. TYPE: Tensor

RETURNS	DESCRIPTION
dict[str, Tensor]	Dictionary containing teacher and student outputs.

Source code in spectre/parametric/distiller.py

def compare_student_teacher(self, X: torch.Tensor) -> dict[str, torch.Tensor]:
    """
    Compare teacher and student outputs on given input.

    Parameters
    ----------
    x : torch.Tensor
        Input data for comparison.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary containing teacher and student outputs.
    """
    self.eval()
    with torch.no_grad():
        s_batch = self.student(X)
        t_batch = self.teacher(X)

    return {
        "teacher_output": t_batch,
        "student_output": s_batch,
        "mse_difference": torch.nn.functional.mse_loss(s_batch, t_batch),
        "relative_error": torch.norm(s_batch - t_batch) / torch.norm(t_batch),
    }