Knowledge Distillation

La distillazione della conoscenza è una tecnica di compressione dei modelli in cui una piccola rete neurale, detta «student», viene addestrata a riprodurre il comportamento di un modello «teacher» più grande e accurato. Lo student apprende le probabilità di output o le caratteristiche interne del teacher, oltre alle etichette di riferimento o al loro posto; questi target più ricchi gli permettono di raggiungere un'accuratezza maggiore rispetto alla stessa architettura addestrata solo sulle etichette. Il risultato è un modello abbastanza rapido e compatto da funzionare su telefoni, schede embedded e altro hardware edge AI, mantenendo gran parte della qualità del teacher. Nel 2006, Buciluǎ, Caruana e Niculescu-Mizil mostrarono che la conoscenza di un ensemble poteva essere compressa in un singolo modello; nel 2015, Hinton, Vinyals e Dean chiamarono questo approccio distillazione della conoscenza e introdussero i target ammorbiditi tramite temperatura, ancora usati dalla maggior parte dei metodi.

Come funziona la distillazione della conoscenza#

Un'etichetta indica solo che un'immagine mostra un lupo. L'output softmax di un teacher addestrato offre più informazioni: l'immagine assomiglia un po' a un husky, un po' a un gatto e per niente a un'auto. Queste probabilità relative delle classi errate, spesso chiamate conoscenza oscura, raggiungono lo student sotto forma di target soft. Dividere i logit per una temperatura T maggiore di 1 prima della softmax appiattisce la distribuzione, dando alle probabilità più basse un peso sufficiente per l'apprendimento. La loss di distillazione è la divergenza di Kullback-Leibler tra le distribuzioni ammorbidite del teacher e dello student, moltiplicata per T² affinché i gradienti mantengano dimensioni comparabili a quelli della normale loss sulle etichette:

import torch
import torch.nn.functional as F

# Logits for one wolf photo over the classes [wolf, husky, cat, car]
teacher_logits = torch.tensor([[6.0, 4.0, 1.0, -2.0]])
student_logits = torch.tensor([[4.0, 1.0, 1.5, 0.0]])

# Dividing logits by a temperature T > 1 softens the teacher's distribution
for T in (1.0, 4.0):
    probs = F.softmax(teacher_logits / T, dim=1)
    print(f"T={T}: teacher probabilities {[round(p, 3) for p in probs[0].tolist()]}")

# Distillation loss: KL divergence between the softened distributions, scaled by T^2
T = 4.0
soft_targets = F.softmax(teacher_logits / T, dim=1)
student_log_probs = F.log_softmax(student_logits / T, dim=1)
loss = F.kl_div(student_log_probs, soft_targets, reduction="batchmean") * T**2
print(f"Distillation loss at T={T}: {loss.item():.3f}")

Con T=1, il teacher assegna alla classe lupo il valore 0.875 e soltanto 0.006 a gatto e 0.0 ad auto. Con T=4, gli stessi logit diventano 0.493, 0.299, 0.141 e 0.067: lo student capisce così che un lupo assomiglia più a un husky che a un gatto, e più a un gatto che a un'auto. Lo student minimizza questa loss, qui pari a 1.244, di solito come somma ponderata con la propria loss del task sulle etichette di riferimento. Il teacher rimane congelato per tutto il processo ed esegue solo passaggi forward:

Tipi di distillazione#

• Basata sulle risposte: lo student imita gli output finali del teacher, come nel metodo dei target soft descritto sopra, e non ha bisogno di altro dal teacher oltre ai suoi logit. DistilBERT ha combinato questo approccio con la modellazione del linguaggio e le loss degli stati nascosti durante il pretraining, rendendo BERT più piccolo del 40% e più veloce del 60%, mantenendo il 97% delle sue capacità di comprensione del linguaggio.
• Basata sulle caratteristiche: lo student imita le mappe delle caratteristiche intermedie del teacher, usando un piccolo proiettore appreso per compensare eventuali differenze nel numero di canali. FitNets introdusse questi «suggerimenti» e addestrò uno student con circa 10.4 volte meno parametri del teacher, superandolo su CIFAR-10.
• A livello di sequenza: un teacher basato su un modello linguistico genera testo e lo student viene sottoposto a fine-tuning su questi output. DeepSeek ha eseguito il fine-tuning dei suoi modelli distillati R1, basati su Qwen2.5 e Llama 3, usando 800,000 campioni selezionati con DeepSeek-R1.

La distillazione basata sulle caratteristiche è adatta ai rilevatori, i cui output comprendono migliaia di riquadri e punteggi anziché un'unica distribuzione di classi. Ultralytics YOLO usa questo approccio quando il training riceve un teacher distill_model: un modello YOLO26 più grande e congelato fornisce le tre mappe delle caratteristiche del neck che alimentano la sua testa di rilevamento, mentre lo student le imita tramite proiettori di convoluzione 1×1, con una loss L2 ponderata per la confidenza del teacher e scalata da dis (6.0 per impostazione predefinita). Su COCO val2017, il modello nano YOLO26 distillato raggiunge 41.5 mAP_50-95, contro 40.9 dello stesso modello addestrato senza teacher, a parità di costo di inferenza.

from ultralytics import YOLO

# Student: il modello YOLO26 nano compatto che verrà distribuito
model = YOLO("yolo26n.pt")

# Teacher: un checkpoint YOLO26 più grande che guida il training e viene eliminato al termine
results = model.train(data="coco8.yaml", epochs=5, distill_model="yolo26s.pt")

Distillazione e termini correlati#

• Transfer learning: riutilizza i pesi del modello come punto di partenza per un nuovo task. La distillazione addestra un'architettura diversa, di solito più piccola, a imitare un modello sullo stesso task.
• Potatura del modello: rimuove pesi, canali o livelli da una rete addestrata. Spesso i modelli potati vengono sottoposti a fine-tuning con la distillazione del modello originale per recuperare accuratezza.
• Quantizzazione del modello: memorizza pesi e attivazioni con una precisione inferiore, ad esempio INT8. Riduce le dimensioni della stessa architettura e si combina con la distillazione, che invece seleziona un'architettura più piccola.
• Ensemble di modelli: calcola la media degli output di diversi modelli in fase di inferenza. Distillare un ensemble in un unico student fu la motivazione originale della tecnica.
• Distillazione dei dataset: nonostante il nome simile, comprime un dataset invece di un modello. Wang et al. hanno compresso le 60.000 immagini di addestramento di MNIST in 10 immagini sintetiche, una per classe, che portano una rete con inizializzazione fissa vicino alla sua precisione originale in pochi passaggi di gradiente.