01: Distillation of Vision Models on CIFAR10¶

This notebook will illustrate how to do the following:

• Use HintonLoss for knowledge distillation at the logits level
• Use the FeatureLoss for knowledge distillation at intermediate levels
• Combine these 2 losses into one with CombinedLoss

The use case example would be the CIFAR10, a benchmark dataset of 60,000 color images sized 32 × 32 pixels, divided evenly into 10 distinct object classes. The goal here is not to reach state-of-the-art performance but to show how to use shrinkai effectively for knowledge distillation.

Setup¶

import torch
import tutorial_utils
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

from shrinkai.distillation import Distiller

import torch import tutorial_utils from torch.utils.data import DataLoader, Subset from torchvision import datasets, transforms from shrinkai.distillation import Distiller

teacher = torch.hub.load(
    "chenyaofo/pytorch-cifar-models",
    "cifar10_vgg19_bn",
    pretrained=True,
    verbose=False,
)

student = torch.hub.load(
    "chenyaofo/pytorch-cifar-models",
    "cifar10_resnet20",
    pretrained=False,
    verbose=False,
)

teacher = torch.hub.load( "chenyaofo/pytorch-cifar-models", "cifar10_vgg19_bn", pretrained=True, verbose=False, ) student = torch.hub.load( "chenyaofo/pytorch-cifar-models", "cifar10_resnet20", pretrained=False, verbose=False, )

def build_cifar_dataloaders(
    batch_size: int = 64,
    subset_size: int | None = 3000,
) -> tuple[DataLoader, DataLoader, torch.Tensor]:
    transform_train = transforms.Compose(
        [
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
        ]
    )

    transform_val = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
        ]
    )

    train_dataset = datasets.CIFAR10(
        root="./data", train=True, download=True, transform=transform_train
    )
    val_dataset = datasets.CIFAR10(
        root="./data", train=False, download=True, transform=transform_val
    )

    if subset_size is not None:
        train_dataset = Subset(train_dataset, range(min(subset_size, len(train_dataset))))
        val_dataset = Subset(val_dataset, range(min(subset_size // 4, len(val_dataset))))

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2)

    sample_input = torch.randn(batch_size, 3, 32, 32)

    return train_loader, val_loader, sample_input

def build_cifar_dataloaders( batch_size: int = 64, subset_size: int | None = 3000, ) -> tuple[DataLoader, DataLoader, torch.Tensor]: transform_train = transforms.Compose( [ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ] ) transform_val = transforms.Compose( [ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ] ) train_dataset = datasets.CIFAR10( root="./data", train=True, download=True, transform=transform_train ) val_dataset = datasets.CIFAR10( root="./data", train=False, download=True, transform=transform_val ) if subset_size is not None: train_dataset = Subset(train_dataset, range(min(subset_size, len(train_dataset)))) val_dataset = Subset(val_dataset, range(min(subset_size // 4, len(val_dataset)))) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2) sample_input = torch.randn(batch_size, 3, 32, 32) return train_loader, val_loader, sample_input

train_loader, val_loader, sample_input = build_cifar_dataloaders(batch_size=128, subset_size=None)

train_loader, val_loader, sample_input = build_cifar_dataloaders(batch_size=128, subset_size=None)

To distill models using shrinkai, you only need to use a Distiller with your 2 models, training parameters and a BaseDistillationLoss object. The package provides many Loss functions, but you can create your own (see Tutorial 05).

Hinton Loss¶

from shrinkai.distillation.losses import HintonLoss

from shrinkai.distillation.losses import HintonLoss

Hinton Loss is the most common Loss for knowledge distillation, and was introduced in the paper Distilling the Knowledge in a Neural Network in 2015:

$$L_{hinton}\left(z_{student}, z_{teacher}, y_{true}\right) = \left(1 - \alpha\right) L_{CE}\left(z_{student}, y_{true}\right) + \alpha \cdot T^2 \cdot L_{KL}\left(\text{softmax}\left(\frac{z_{student}}{T}\right), \text{softmax}\left(\frac{z_{teacher}}{T}\right)\right)$$

with:

• $z_{student}$: The logits output by the student model.
• $z_{teacher}$: The logits output by the larger, pre-trained teacher model.
• $y_{true}$: The ground-truth labels from the actual dataset.
• $L_{CE}$: The standard Cross-Entropy loss function.
• $L_{KL}$: The Kullback-Leibler Divergence loss. Measures the distance between the student's soft probability distribution and the teacher's soft probability distribution. It is sometimes referred as $D_{KL}$.
• $\alpha$: A hyperparameter between 0 and 1 that balances the two losses. Bigger $\alpha$ means teacher's logits are more important than ground truth.
• $T$ (temperature): A hyperparameter used to divide the logits before the softmax function. A higher temperature "softens" the probabilities, revealing the secondary class relationships ("dark knowledge") learned by the teacher. Dividing by $T$ shrinks the gradients by $1/T^2$, so the distillation loss must be multiplied by $T^2$ to keep it balanced with the standard cross-entropy loss.

This Loss function allows the student to learn from the teacher responses but also from the ground truth (depending on $\alpha$). However, it limits the distillation strictly to the final logits, it does not allow the student to learn from the teacher's intermediate hidden states or feature maps.

hinton_loss = HintonLoss(temperature=4.0, alpha=0.7)

hinton_loss = HintonLoss(temperature=4.0, alpha=0.7)

distiller = Distiller(
    teacher=teacher,
    student=student,
    criterion=hinton_loss,
    optimizer="adamw",
    lr=5e-4,
    weight_decay=1e-4,
    device="auto",
)

distiller = Distiller( teacher=teacher, student=student, criterion=hinton_loss, optimizer="adamw", lr=5e-4, weight_decay=1e-4, device="auto", )

training_history = distiller.fit(
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    epochs=15,
)

training_history = distiller.fit( train_dataloader=train_loader, val_dataloader=val_loader, epochs=15, )

Epoch [01/15] Train Loss: 8.1248 - Train Acc: 38.98% | Val Loss: 6.2871 - Val Acc: 48.23%

Epoch [02/15] Train Loss: 6.2774 - Train Acc: 54.23% | Val Loss: 6.1256 - Val Acc: 52.41%

Epoch [03/15] Train Loss: 5.3766 - Train Acc: 61.45% | Val Loss: 4.7000 - Val Acc: 61.58%

Epoch [04/15] Train Loss: 4.7910 - Train Acc: 65.71% | Val Loss: 4.2208 - Val Acc: 66.35%

Epoch [05/15] Train Loss: 4.3816 - Train Acc: 68.75% | Val Loss: 4.2195 - Val Acc: 66.24%

Epoch [06/15] Train Loss: 3.9922 - Train Acc: 71.83% | Val Loss: 3.8216 - Val Acc: 69.99%

Epoch [07/15] Train Loss: 3.7109 - Train Acc: 73.63% | Val Loss: 3.2321 - Val Acc: 73.81%

Epoch [08/15] Train Loss: 3.4792 - Train Acc: 75.38% | Val Loss: 3.0538 - Val Acc: 74.55%

Epoch [09/15] Train Loss: 3.2924 - Train Acc: 76.94% | Val Loss: 3.1627 - Val Acc: 74.65%

Epoch [10/15] Train Loss: 3.1259 - Train Acc: 77.96% | Val Loss: 3.0203 - Val Acc: 75.37%

Epoch [11/15] Train Loss: 2.9972 - Train Acc: 78.77% | Val Loss: 2.5770 - Val Acc: 78.55%

Epoch [12/15] Train Loss: 2.8949 - Train Acc: 79.77% | Val Loss: 2.5373 - Val Acc: 78.97%

Epoch [13/15] Train Loss: 2.7803 - Train Acc: 80.52% | Val Loss: 2.4638 - Val Acc: 79.53%

Epoch [14/15] Train Loss: 2.6700 - Train Acc: 81.32% | Val Loss: 2.4322 - Val Acc: 79.83%

Epoch [15/15] Train Loss: 2.5857 - Train Acc: 81.98% | Val Loss: 2.3659 - Val Acc: 80.08%

tutorial_utils.plot_history(training_history)

tutorial_utils.plot_history(training_history)

The student model has only been trained for a small number of epochs here, to keep the tutorial short. It seems that it has not reached the best performance yet, but it is not the goal here. As the training is over, a small report can be displayed and the student model can be saved:

report = distiller.benchmark(
    sample_input=sample_input,
    teacher_name="VGG19-BN (CIFAR-10)",
    student_name="ResNet-20 (Distilled)",
    val_dataloader=val_loader,
)

report.show()

report = distiller.benchmark( sample_input=sample_input, teacher_name="VGG19-BN (CIFAR-10)", student_name="ResNet-20 (Distilled)", val_dataloader=val_loader, ) report.show()

                                         Distillation Benchmark Report                                          
┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Metric            ┃ Teacher (VGG19-BN (CIFAR-10)) ┃ Student (ResNet-20 (Distilled)) ┃     Gain / Compression ┃
┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━┩
│ Parameters        │                       20.57 M │                          0.27 M │                 -98.7% │
│ Model Size (Disk) │                      78.53 MB │                         1.08 MB │ -98.6% (72.7x smaller) │
│ Latency / Sample  │                       0.20 ms │                         0.04 ms │            4.6x faster │
│ Throughput (FPS)  │                  4943.4 img/s │                   22851.3 img/s │    +4.6x (22851.3 FPS) │
│ Accuracy          │                        93.39% │                          80.08% │         85.7% retained │
└───────────────────┴───────────────────────────────┴─────────────────────────────────┴────────────────────────┘

You can save your model directly from the Distiller, but more export features will be shown in Tutorial 06.

output_path = "../../models/tutorial_01_hinton.pt"
distiller.save_student(output_path)

output_path = "../../models/tutorial_01_hinton.pt" distiller.save_student(output_path)

Feature Loss¶

from shrinkai.adapters import FeatureExtractor, FeatureProjector
from shrinkai.distillation.losses import FeatureLoss, ProjectedFeatureLoss

from shrinkai.adapters import FeatureExtractor, FeatureProjector from shrinkai.distillation.losses import FeatureLoss, ProjectedFeatureLoss

While Hinton Loss focuses on the final output, Feature-Based Distillation forces the student to mimic the internal reasoning process of the teacher. Introduced in the paper FitNets: Hints for Thin Deep Nets (2014), this approach aligns the intermediate hidden states (or feature maps) of the two models.

$$L_{feature}\left(f_{student}, f_{teacher}\right) = \text{Distance}\left(P(f_{student}), f_{teacher}\right)$$

with:

• $f_{student}$: The intermediate feature map extracted from a specific layer of the student model.
• $f_{teacher}$: The intermediate feature map extracted from a specific layer of the teacher model.
• $P$ (Projector): A transformation function (usually a 1x1 Convolution or a Linear layer) applied to the student's features. Because the student is often smaller, its feature maps have fewer channels than the teacher's. The projector aligns the student's dimensions with the teacher's dimensions.
• $\text{Distance}$: The metric used to compute the difference between the two feature maps (e.g., Mean Squared Error, L1, or Cosine Similarity).

This Loss function allows the student to learn hierarchical representations and rich spatial or semantic features directly from the teacher. It is especially powerful when the teacher is much deeper than the student. However, performing Knowledge Distillation at intermediate levels (hidden features, attention maps) is far more challenging than standard logit distillation. Because the student and teacher often have completely different architectures, their internal states are not only difficult to extract, but they also have mismatched channel dimensions and layer names.

To solve this elegantly, shrinkai provides 3 dedicated classes:

• FeatureExtractor: Dynamically intercepts intermediate activations during the forward pass (without modifying the original models' source code).
• FeatureProjector: Aligns the student's feature dimensions with the teacher's using projection layers (the $P$ function in the equation above).
• ProjectedFeatureLoss: A seamless bridge that orchestrates the extractor, the projector, and your chosen feature loss into a single criterion ready for training.

To get your neural network architecture, you can use the following lines.

print(dict(teacher.named_modules()).keys())

print(dict(teacher.named_modules()).keys())

dict_keys(['', 'features', 'features.0', 'features.1', 'features.2', 'features.3', 'features.4', 'features.5', 'features.6', 'features.7', 'features.8', 'features.9', 'features.10', 'features.11', 'features.12', 'features.13', 'features.14', 'features.15', 'features.16', 'features.17', 'features.18', 'features.19', 'features.20', 'features.21', 'features.22', 'features.23', 'features.24', 'features.25', 'features.26', 'features.27', 'features.28', 'features.29', 'features.30', 'features.31', 'features.32', 'features.33', 'features.34', 'features.35', 'features.36', 'features.37', 'features.38', 'features.39', 'features.40', 'features.41', 'features.42', 'features.43', 'features.44', 'features.45', 'features.46', 'features.47', 'features.48', 'features.49', 'features.50', 'features.51', 'features.52', 'classifier', 'classifier.0', 'classifier.1', 'classifier.2', 'classifier.3', 'classifier.4', 'classifier.5', 'classifier.6'])

print(dict(student.named_modules()).keys())

print(dict(student.named_modules()).keys())

dict_keys(['', 'conv1', 'bn1', 'relu', 'layer1', 'layer1.0', 'layer1.0.conv1', 'layer1.0.bn1', 'layer1.0.relu', 'layer1.0.conv2', 'layer1.0.bn2', 'layer1.1', 'layer1.1.conv1', 'layer1.1.bn1', 'layer1.1.relu', 'layer1.1.conv2', 'layer1.1.bn2', 'layer1.2', 'layer1.2.conv1', 'layer1.2.bn1', 'layer1.2.relu', 'layer1.2.conv2', 'layer1.2.bn2', 'layer2', 'layer2.0', 'layer2.0.conv1', 'layer2.0.bn1', 'layer2.0.relu', 'layer2.0.conv2', 'layer2.0.bn2', 'layer2.0.downsample', 'layer2.0.downsample.0', 'layer2.0.downsample.1', 'layer2.1', 'layer2.1.conv1', 'layer2.1.bn1', 'layer2.1.relu', 'layer2.1.conv2', 'layer2.1.bn2', 'layer2.2', 'layer2.2.conv1', 'layer2.2.bn1', 'layer2.2.relu', 'layer2.2.conv2', 'layer2.2.bn2', 'layer3', 'layer3.0', 'layer3.0.conv1', 'layer3.0.bn1', 'layer3.0.relu', 'layer3.0.conv2', 'layer3.0.bn2', 'layer3.0.downsample', 'layer3.0.downsample.0', 'layer3.0.downsample.1', 'layer3.1', 'layer3.1.conv1', 'layer3.1.bn1', 'layer3.1.relu', 'layer3.1.conv2', 'layer3.1.bn2', 'layer3.2', 'layer3.2.conv1', 'layer3.2.bn1', 'layer3.2.relu', 'layer3.2.conv2', 'layer3.2.bn2', 'avgpool', 'fc'])

Before starting the distillation, some configs must be done.

# Make sure these 2 dictionnaries have the same values

teacher_layers = {
    "features.13": "stage_1",
    "features.26": "stage_2",
    "features.39": "stage_3",
}
student_layers = {
    "layer1": "stage_1",
    "layer2": "stage_2",
    "layer3": "stage_3",
}

# Make sure these 2 dictionnaries have the same values teacher_layers = { "features.13": "stage_1", "features.26": "stage_2", "features.39": "stage_3", } student_layers = { "layer1": "stage_1", "layer2": "stage_2", "layer3": "stage_3", }

wrapped_teacher = FeatureExtractor(teacher, target_layers=teacher_layers)
wrapped_student = FeatureExtractor(student, target_layers=student_layers)

wrapped_teacher = FeatureExtractor(teacher, target_layers=teacher_layers) wrapped_student = FeatureExtractor(student, target_layers=student_layers)

# Make sure the keys are the values of the teacher/student_layers
# and that the dimension are right

proj_config = {
    "stage_1": {
        "in_channels": 16,
        "out_channels": 128,
        "type": "conv",
        "use_norm": True,
    },
    "stage_2": {
        "in_channels": 32,
        "out_channels": 256,
        "type": "conv",
        "use_norm": True,
    },
    "stage_3": {
        "in_channels": 64,
        "out_channels": 512,
        "type": "conv",
        "use_norm": True,
    },
}

# Make sure the keys are the values of the teacher/student_layers # and that the dimension are right proj_config = { "stage_1": { "in_channels": 16, "out_channels": 128, "type": "conv", "use_norm": True, }, "stage_2": { "in_channels": 32, "out_channels": 256, "type": "conv", "use_norm": True, }, "stage_3": { "in_channels": 64, "out_channels": 512, "type": "conv", "use_norm": True, }, }

projector = FeatureProjector(proj_config)

feature_loss = ProjectedFeatureLoss(
    projector=projector, feature_loss=FeatureLoss(loss_type="mse", normalize=True)
)

projector = FeatureProjector(proj_config) feature_loss = ProjectedFeatureLoss( projector=projector, feature_loss=FeatureLoss(loss_type="mse", normalize=True) )

Now, everything is ready for the distillation.

distiller = Distiller(
    teacher=wrapped_teacher,
    student=wrapped_student,
    criterion=feature_loss,
    optimizer="adamw",
    lr=5e-4,
    weight_decay=1e-4,
    device="auto",
)

distiller = Distiller( teacher=wrapped_teacher, student=wrapped_student, criterion=feature_loss, optimizer="adamw", lr=5e-4, weight_decay=1e-4, device="auto", )

training_history = distiller.fit(
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    epochs=15,
)

training_history = distiller.fit( train_dataloader=train_loader, val_dataloader=val_loader, epochs=15, )

Epoch [01/15] Train Loss: 0.0003 - Train Acc: 26.96% | Val Loss: 0.0002 - Val Acc: 27.07%

Epoch [02/15] Train Loss: 0.0002 - Train Acc: 28.04% | Val Loss: 0.0002 - Val Acc: 37.40%

Epoch [03/15] Train Loss: 0.0002 - Train Acc: 37.00% | Val Loss: 0.0002 - Val Acc: 43.82%

Epoch [04/15] Train Loss: 0.0001 - Train Acc: 42.43% | Val Loss: 0.0002 - Val Acc: 46.00%

Epoch [05/15] Train Loss: 0.0001 - Train Acc: 44.63% | Val Loss: 0.0002 - Val Acc: 48.64%

Epoch [06/15] Train Loss: 0.0001 - Train Acc: 46.29% | Val Loss: 0.0001 - Val Acc: 54.54%

Epoch [07/15] Train Loss: 0.0001 - Train Acc: 49.64% | Val Loss: 0.0001 - Val Acc: 51.96%

Epoch [08/15] Train Loss: 0.0001 - Train Acc: 52.35% | Val Loss: 0.0001 - Val Acc: 58.46%

Epoch [09/15] Train Loss: 0.0001 - Train Acc: 53.74% | Val Loss: 0.0001 - Val Acc: 58.61%

Epoch [10/15] Train Loss: 0.0001 - Train Acc: 55.33% | Val Loss: 0.0001 - Val Acc: 57.56%

Epoch [11/15] Train Loss: 0.0001 - Train Acc: 55.85% | Val Loss: 0.0001 - Val Acc: 57.73%

Epoch [12/15] Train Loss: 0.0001 - Train Acc: 56.17% | Val Loss: 0.0001 - Val Acc: 61.60%

Epoch [13/15] Train Loss: 0.0001 - Train Acc: 57.14% | Val Loss: 0.0001 - Val Acc: 62.68%

Epoch [14/15] Train Loss: 0.0001 - Train Acc: 57.47% | Val Loss: 0.0001 - Val Acc: 60.34%

Epoch [15/15] Train Loss: 0.0001 - Train Acc: 57.87% | Val Loss: 0.0001 - Val Acc: 61.12%

tutorial_utils.plot_history(training_history)

tutorial_utils.plot_history(training_history)

As observed, the accuracy is not stable and remains virtually constant. One could have expected this behavior because FeatureLoss exclusively optimizes the intermediate representations (hidden states) rather than the final classification head. As a consequence, the predictive layer remains untrained, resulting in an accuracy of approximately 10%, the exact baseline for random guessing on the CIFAR-10 dataset.

However, FeatureLoss is highly valuable: it forces the student model to capture the deep, representational logic of the teacher. While this structural learning is not directly reflected in standard classification metrics, it builds a robust foundation. To leverage these rich intermediate features while actually training the model to predict correctly, combining feature-based and logit-based losses becomes essential.

By the way, the feature_analysis method from Distiller, like benchmark, gives important metrics to evaluate how well the student has learned from the teacher for each feature. So far, there are 3 metrics:

• CKA (Centered Kernel Alignment): Measures the absolute similarity between the two feature spaces using their Gram matrices ($K$ for teacher, $L$ for student, just like in the original paper). $$CKA(K, L) = \frac{HSIC(K, L)}{\sqrt{HSIC(K, K) HSIC(L, L)}}$$

  • HSIC is the Hilbert-Schmidt Independence Criterion.
  • Range: $[0, 1]$
  • Interpretation: A score close to 1 indicates that both models extract highly correlated linear features.

• RSA (Representational Similarity Analysis): Measures the topological geometry of the features by comparing the pairwise distances between samples in a batch using Pearson correlation ($\rho$). Let $R_S$ and $R_T$ be the Representational Dissimilarity Matrices (RDMs) of the student and teacher. $$RSA = \frac{\text{cov}(R_S, R_T)}{\sigma_{R_S} \sigma_{R_T}}$$

  • Range: $[-1, 1]$
  • Interpretation: A high positive score means that if the teacher considers two images to be visually/semantically similar, the student does too.

• Spatial Attention: Collapses each feature map's channel dimension into a single spatial saliency map (Zagoruyko & Komodakis's "attention transfer", $A = \sum_c |f_c|$), then compares the two saliency maps via Cosine Similarity. $$Spatial = \frac{\sum (A_{student} \odot A_{teacher})}{\|A_{student}\| \|A_{teacher}\|}$$

  • Range: $[0, 1]$
  • Interpretation: It answers the question: "Are the two models looking at the same regions of the image?". A score close to 1 means that when the teacher's activations concentrate over a specific region (e.g., the object's silhouette), the student's activations concentrate over that same region too, even though the two networks compute entirely different features there.

Note: this last metric works on any 4D tensor, not only CNN activation maps. In the NLP tutorials, it is applied directly to real self-attention probability matrices ([Batch, Heads, Seq_Len, Seq_Len]) instead, but it then compares attention weight patterns between teacher and student rather than spatial feature saliency.

analysis_report = distiller.feature_analysis(val_loader, metrics=["cka", "rsa", "attention"])
analysis_report.show()

analysis_report = distiller.feature_analysis(val_loader, metrics=["cka", "rsa", "attention"]) analysis_report.show()

          Representation Alignment Report (Feature Distillation)          
┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓
┃ Layer / Stage Alias ┃ CKA (Linear) ┃ RSA (Pearson) ┃ Spatial Attention ┃
┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩
│ stage_3             │        0.358 │         0.562 │             0.992 │
│ stage_2             │        0.718 │         0.571 │             0.949 │
│ stage_1             │        0.919 │         0.447 │             0.936 │
└─────────────────────┴──────────────┴───────────────┴───────────────────┘

Hybrid Loss¶

from shrinkai.distillation.losses import HybridLoss

from shrinkai.distillation.losses import HybridLoss

HybridLoss aims to use a Loss on logits and a Loss on hidden states. Here are the 2 ways to use it:

• Additive (convex_weighting=False): $L = L_{logits} + \text{weight} \times L_{features}$
• Convex (convex_weighting=True): $L = (1 - \text{weight}) \times L_{logits} + \text{weight} \times L_{features}$

For this tutorial, let's use the Additive mode and just add the 2 previous Losses with a feature_weight equals to 1000 in order to have 2 terms with the same scale, according to the previous trainings.

hybrid_loss = HybridLoss(
    logit_loss=hinton_loss,
    feature_loss=feature_loss,
    feature_weight=1000,
    convex_weighting=False,
)

hybrid_loss = HybridLoss( logit_loss=hinton_loss, feature_loss=feature_loss, feature_weight=1000, convex_weighting=False, )

distiller = Distiller(
    teacher=wrapped_teacher,
    student=wrapped_student,
    criterion=hybrid_loss,
    optimizer="adamw",
    lr=5e-4,
    weight_decay=1e-4,
    device="auto",
)

distiller = Distiller( teacher=wrapped_teacher, student=wrapped_student, criterion=hybrid_loss, optimizer="adamw", lr=5e-4, weight_decay=1e-4, device="auto", )

training_history = distiller.fit(
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    epochs=15,
)

training_history = distiller.fit( train_dataloader=train_loader, val_dataloader=val_loader, epochs=15, )

Epoch [01/15] Train Loss: 2.8174 - Train Acc: 82.57% | Val Loss: 2.2414 - Val Acc: 81.91%

Epoch [02/15] Train Loss: 2.3564 - Train Acc: 84.28% | Val Loss: 2.0166 - Val Acc: 82.91%

Epoch [03/15] Train Loss: 2.2386 - Train Acc: 85.16% | Val Loss: 2.0094 - Val Acc: 83.23%

Epoch [04/15] Train Loss: 2.1564 - Train Acc: 85.54% | Val Loss: 1.9953 - Val Acc: 83.66%

Epoch [05/15] Train Loss: 2.0891 - Train Acc: 86.06% | Val Loss: 2.0047 - Val Acc: 83.14%

Epoch [06/15] Train Loss: 2.0417 - Train Acc: 86.32% | Val Loss: 1.9160 - Val Acc: 84.09%

Epoch [07/15] Train Loss: 1.9887 - Train Acc: 86.68% | Val Loss: 1.7940 - Val Acc: 85.05%

Epoch [08/15] Train Loss: 1.9349 - Train Acc: 87.11% | Val Loss: 1.7186 - Val Acc: 85.28%

Epoch [09/15] Train Loss: 1.8727 - Train Acc: 87.44% | Val Loss: 1.8424 - Val Acc: 84.25%

Epoch [10/15] Train Loss: 1.8713 - Train Acc: 87.58% | Val Loss: 1.9761 - Val Acc: 84.04%

Epoch [11/15] Train Loss: 1.8172 - Train Acc: 87.79% | Val Loss: 1.7358 - Val Acc: 85.06%

Epoch [12/15] Train Loss: 1.7719 - Train Acc: 88.13% | Val Loss: 1.9346 - Val Acc: 83.87%

Epoch [13/15] Train Loss: 1.7316 - Train Acc: 88.49% | Val Loss: 1.7782 - Val Acc: 85.41%

Epoch [14/15] Train Loss: 1.6922 - Train Acc: 88.72% | Val Loss: 1.6366 - Val Acc: 86.33%

Epoch [15/15] Train Loss: 1.6593 - Train Acc: 89.18% | Val Loss: 1.5919 - Val Acc: 86.47%

tutorial_utils.plot_history(training_history)

tutorial_utils.plot_history(training_history)

distiller.benchmark(
    sample_input=sample_input,
    teacher_name="VGG19-BN (CIFAR-10)",
    student_name="ResNet-20 (Distilled)",
    val_dataloader=val_loader,
).show()

distiller.benchmark( sample_input=sample_input, teacher_name="VGG19-BN (CIFAR-10)", student_name="ResNet-20 (Distilled)", val_dataloader=val_loader, ).show()

                                         Distillation Benchmark Report                                          
┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Metric            ┃ Teacher (VGG19-BN (CIFAR-10)) ┃ Student (ResNet-20 (Distilled)) ┃     Gain / Compression ┃
┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━┩
│ Parameters        │                       20.57 M │                          0.27 M │                 -98.7% │
│ Model Size (Disk) │                      78.53 MB │                         1.08 MB │ -98.6% (72.7x smaller) │
│ Latency / Sample  │                       0.27 ms │                         0.05 ms │            5.3x faster │
│ Throughput (FPS)  │                  3640.4 img/s │                   19175.8 img/s │    +5.3x (19175.8 FPS) │
│ Accuracy          │                        93.39% │                          86.47% │         92.6% retained │
└───────────────────┴───────────────────────────────┴─────────────────────────────────┴────────────────────────┘

distiller.feature_analysis(val_loader, metrics=["cka", "rsa", "attention"]).show()

distiller.feature_analysis(val_loader, metrics=["cka", "rsa", "attention"]).show()

          Representation Alignment Report (Feature Distillation)          
┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓
┃ Layer / Stage Alias ┃ CKA (Linear) ┃ RSA (Pearson) ┃ Spatial Attention ┃
┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩
│ stage_3             │        0.816 │         0.748 │             0.988 │
│ stage_2             │        0.701 │         0.557 │             0.939 │
│ stage_1             │        0.927 │         0.453 │             0.937 │
└─────────────────────┴──────────────┴───────────────┴───────────────────┘

In this example, combining a logit loss with a feature loss substantially improved the alignment between the student's and teacher's intermediate representations (see the CKA/RSA jump above, compared to the pure FeatureLoss run), while keeping the classification accuracy roughly in the same range as the Hinton-only baseline from earlier. The feature term does not directly optimize accuracy, but it shapes the internal representations the classifier head builds on, which matters most on tasks where those representations get reused (transfer learning, further compression, etc.).

References¶

1. Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the Knowledge in a Neural Network. arXiv:1503.02531.
   https://arxiv.org/abs/1503.02531

2. Romero, A., Ballas, N., Ebrahimi Kahou, S., Chassang, A., Gatta, C., & Bengio, Y. (2015). FitNets: Hints for Thin Deep Nets. ICLR. arXiv:1412.6550.
   https://arxiv.org/abs/1412.6550

3. Zagoruyko, S., & Komodakis, N. (2017). Paying More Attention to Attention: Improving the Performance of Convolutional Neural Networks via Attention Transfer. ICLR. arXiv:1612.03928.
   https://arxiv.org/abs/1612.03928

4. Kornblith, S., Norouzi, M., Lee, H., & Hinton, G. (2019). Similarity of Neural Network Representations Revisited. ICML, PMLR 97, 3519–3529. arXiv:1905.00414.
   https://arxiv.org/abs/1905.00414

5. Gretton, A., Bousquet, O., Smola, A., & Schölkopf, B. (2005). Measuring Statistical Dependence with Hilbert-Schmidt Norms. ALT 2005, pp. 63–77.
   DOI: 10.1007/11564089_7

6. Kriegeskorte, N., Mur, M., & Bandettini, P. (2008). Representational Similarity Analysis – Connecting the Branches of Systems Neuroscience. Frontiers in Systems Neuroscience, 2, 4.
   DOI: 10.3389/neuro.06.004.2008

7. Krizhevsky, A. (2009). Learning Multiple Layers of Features from Tiny Images. Technical Report, University of Toronto.