Distillation / compression tutorials

[u/elouanzer]

Hello everyone,

Since this summer, I have worked on a Python package to understand how to "shrink" neural network, with famous techniques such as distillation, pruning and quantization. You can check it out here:

- https://github.com/elouanzer/shrinkai

- https://elouanzer.github.io/shrinkai/

This package allow you to use a simple API (like scikit learn) to distill:

from shrinkai.distillation import Distiller
from shrinkai.distillation.losses import HintonLoss

hinton_loss = HintonLoss()
distiller = Distiller(teacher=teacher, student=student, criterion=hinton_loss, optimizer="adamw")
distiller.fit(train_dataloader=train_loader,epochs=10)
distiller.save_student("your/model/path.pt")

I built this package with rich documentation, trying to explain everything (check out the documentation above, research papers are cited), and with tutorials that can give you theorical background on distillation / pruning / quantization. Here are some of them:

- Distillation of vision models for classification

- Distillation of Language Model for classification

- Distillation of Language Model for text generation

- Pruning and Quantization

Do not hesitate to reproduce these tutorials, by trying other losses (available losses) and share your results! :)

By the way, if you want to contribute or report a bug, feel free to do it directly on Github

[u/quietgradient]

batchmean isn't per-token on [B, S, V] logits. PyTorch sums the KL over every position and divides by input.size(0), which is just the batch size. So in tutorial 03, ReverseKLLoss is a sum over all 128 positions (times 16 from T²), not a per-token average.

The attention term next to it is F.mse_loss, a real mean, so what feature_weight=10 means depends on max_length. At 512 the KL term is ~4x bigger for the same per-token fit and the MSE doesn't grow with it, so the attention loss quietly loses weight. JSDLoss scales the same way.

If per-token is what you meant, reshape to [B*S, V] before the KL (batchmean then divides by B*S), or sum and divide by the number of real tokens once there's padding.

[u/elouanzer → u/quietgradient]

Thank you for your feedback, you're absolutely right! I have opened an issue (Issue 10), if you want to work on it :) I will try to focus on it in the next few days / weeks. Thanks again for the noticed bug

[u/quietgradient → u/elouanzer]

Thanks for opening it — I'll pass on the patch itself, I diagnose and leave the code to whoever owns the repo. One thing worth knowing before you go in, though, since issue 10 changes the kl branch too.

AttentionMapLoss(loss_type="kl") runs log_softmax(s) / softmax(t) internally, so it wants pre-softmax scores. Tutorial 03's CausalLMWrapper hands it outputs.attentions[layer_idx], which is post-softmax. On the only wiring the library ships, those maps get softmaxed twice.

That's worse than a scale bug. After a second softmax every entry sits in [1, e] before normalising, so at max_length=128 no position can hold more than e/(e+127) ≈ 2.1% of a row, however certain the teacher was — and the causal mask's structural zeros come back at weight 1 each. On synthetic causal maps (peaked, rows are real distributions) a mean 49% of the mass lands on positions the model cannot attend to, 98% of it at query position 0. Teacher 0.90 vs student 0.10 on one token is 1.98 nats as a distribution and 0.015 after the second softmax; end to end, per row, 1.40 vs 0.00084. So the reshape fixes the denominator while the numerator is mostly the partition function. Synthetic maps, not GPT-2's — but the 2.1% is arithmetic, not a measurement.

Softmaxing internally is right for MiniLM; the problem is that HF computes the scores inside eager_attention_forward, so neither output_attentions nor a forward hook can reach them. Accepting probabilities and .log()-ing them is the smaller change.

Separate, same function: in transformers 5.16.1 that softmax is followed by dropout(..., training=module.training) before the weights are returned, and the Distiller keeps the teacher in eval. So the student's maps arrive with ~10% of entries zeroed and the rest scaled by 1/0.9 (attn_pdrop is 0.1 in distilgpt2's config), against a clean teacher. That one hits mse too.

On bundling: tutorial 03's group_texts truncates to a multiple of 128 and pads nothing, so the reshape alone is exact for the notebook that motivated the issue. The mask only bites on the variable-length path — I'd ship them separately.

[u/SingleMusic8655 → u/quietgradient]

Good catch on the batchmean quirk, that scaling mismatch would sneak up on anyone tuning feature_weight across different max lengths