Cos'è la distillazione della conoscenza?

La KD trae vantaggio dal fatto che le reti neurali artificiali sono "approssimatori universali": con un numero sufficiente di dati di addestramento e un livello nascosto sufficientemente grande, una rete neurale può approssimare qualsiasi funzione con precisione arbitraria.⁵

Nel machine learning convenzionale, la "conoscenza" di un modello addestrato si identifica con i parametri che ha appreso: i pesi (e i bias) variabili, applicati alle diverse operazioni matematiche che avvengono nella rete neurale, che amplificano o riducono l'influenza che una certa parte dell'output della rete ha su un'altra parte. Questa visione della conoscenza rende difficile capire come un modello possa assorbire la conoscenza di un altro modello di dimensioni e struttura diverse.

Invece, Hinton et al. hanno applicato una visione più astratta e flessibile della conoscenza come semplicemente "una mappatura appresa da vettori di input a vettori di output". In altre parole, la KD interpreta la conoscenza di un modello non come i parametri strettamente matematici che apprende durante l'addestramento, ma come il modo in cui è in grado di generalizzare a nuovi dati dopo l'addestramento.

Attraverso questa concezione alternativa della conoscenza, i metodi di distillazione della conoscenza mirano ad addestrare i modelli studente a imitare non solo l'output finale del modello insegnante dato un determinato input, ma anche i processi di pensiero che il modello insegnante adotta per arrivare a quel risultato finale. Concettualmente, questo funziona in modo simile all' instruction tuning attraverso il i chain-of-thought (CoT) prompting, che migliora la qualità delle risposte degli LLM istruendoli su come articolare la loro logica "passo dopo passo".

Nell'apprendimento convenzionale supervisionato o autosupervisionato, una funzione di perdita produce un vettore che rappresenta la divergenza (o perdita) tra gli output del modello e gli output "corretti" (o ground truth) tra diversi input. Regolando i parametri del modello per ridurre al minimo la pendenza (o il gradiente) di questo vettore attraverso un algoritmo di ottimizzazione come la discesa del gradiente, gli output del modello si avvicinano a quelli corretti. Sebbene le fasi di ragionamento del modello siano "importanti" in quanto influenzano il suo output finale, in genere non vengono misurate da una funzione di perdita convenzionale.

La distillazione della conoscenza, al contrario, allena anche il modello studente a imitare il processo di ragionamento del modello insegnante attraverso l'aggiunta di un tipo specializzato di funzione di perdita, chiamata distillation loss, che utilizza fasi di ragionamento distinte come soft target per l'ottimizzazione.

L'output di qualsiasi modello AI può essere inteso come una serie di previsioni: un LLM autoregressivo prevede le parole successive in una sequenza specificata; un modello di computer vision utilizzato per la classificazione delle immagini prevede la categoria di una determinata immagine. Per arrivare a queste previsioni finali, chiamate "hard target" in questo contesto, i modelli di deep learning in genere fanno più previsioni preliminari e utilizzano una funzione softmax per produrre la previsione con la più alta probabilità. Durante l'addestramento, viene utilizzata una funzione di perdita di entropia incrociata per massimizzare la probabilità assegnata all'output corretto e ridurre al minimo la probabilità assegnata agli output errati.

Ad esempio, un modello di classificazione delle immagini prevede la probabilità che un'immagine di input appartenga a ciascuna classe nota che il modello è addestrato a riconoscere, quindi restituisce la classe con il valore di probabilità più alto. Nel linguaggio matematico del machine learning, queste previsioni individuali per classe sono chiamate logit. Analogamente, un LLM autoregressivo prevede diverse possibilità per ogni parola successiva e (a seconda della sua temperatura ) campiona una di queste possibilità per il suo output.

Nella distillazione della conoscenza, queste previsioni intermedie, i "soft target", generate dal modello insegnante, spesso forniscono i principali dati di addestramento per il modello studente. Le probabilità relative assegnate a queste previsioni preliminari forniscono insight preziosi sul modo in cui il modello insegnante tende a generalizzare. Ad esempio, un modello di classificazione delle immagini ha molte più probabilità di classificare erroneamente l’immagine di una volpe come "cane" piuttosto che come "sandwich". I soft target forniscono quindi molte più informazioni per caso di addestramento rispetto ai soli hard target.

I soft target offrono inoltre maggiore coerenza rispetto agli hard target: la previsione finale di un modello potrebbe in definitiva dipendere da una minuscola differenza tra due valori logit, ma i valori logit stessi hanno una varianza molto inferiore nel gradiente tra ogni esempio di addestramento.

Grazie alla ricchezza e alla stabilità delle informazioni fornite dai soft target, il modello studente può essere addestrato su un minor numero di esempi di addestramento, utilizzando un tasso di apprendimento più elevato, rispetto a quelli utilizzati per addestrare il modello insegnante originale.

Per avvicinare le tendenze di generalizzazione della rete studente a quelle della rete insegnante, la distillazione della conoscenza utilizza in genere due funzioni di perdita. La prima è una funzione di perdita standard che opera in base alla "perdita effettiva", misurando gli output finali del modello studente rispetto alle etichette della ground truth (nell'apprendimento supervisionato) o rispetto al campione di dati originale (nell'apprendimento autosupervisionato). La seconda è la distillation loss, una "soft loss" che misura i soft target del modello studente rispetto a quelli del modello insegnante.

Poiché possono esserci più soft target per ogni esempio di addestramento, la distillation loss misura la differenza tra la distribuzione di probabilità dei soft target della rete insegnante e la distribuzione di probabilità di quelli della rete studente. A questo scopo viene comunemente utilizzata la divergenza di Kullback-Leibler (o "divergenza KL").

Queste diverse forme di conoscenza rientrano generalmente in una delle tre categorie: conoscenza basata sulle risposte, conoscenza basata sulle caratteristiche o conoscenza basata sulle relazioni.

La conoscenza basata sulle risposte, il genere più comune di distillazione della conoscenza, si concentra sul trasferimento di informazioni dal livello di output finale del modello insegnante. In un tipico metodo KD basato sulle risposte, il modello studente viene addestrato a produrre logit che corrispondono alle previsioni del modello insegnante.

Quando i soft target del modello insegnante hanno una bassa entropia, in altre parole, quando le previsioni sono estremamente "sicure", come se un modello di classificazione producesse un logit molto vicino a 1 (che rappresenta la certezza) per una classe e logit che si avvicinano a 0 per tutte le altre, non forniscono molte informazioni. I metodi basati sulle risposte, quindi, utilizzano spesso un'impostazione di temperatura elevata per gli output del modello, che aumenta l'entropia delle previsioni del modello. Ciò garantisce una distribuzione di probabilità più variabile e quindi una maggiore quantità di informazioni da ogni esempio di addestramento.

La conoscenza basata sulle caratteristiche si concentra sulle informazioni trasmesse nei livelli intermedi, o "livelli nascosti", di una rete neurale. È qui che le reti neurali tendono a eseguire l'estrazione delle caratteristiche, ovvero l'identificazione di caratteristiche e pattern distinti dei dati di input rilevanti per l'attività da svolgere.

Ad esempio, nelle reti neurali convoluzionali utilizzate principalmente per attività di computer vision come la segmentazione delle immagini, ogni livello nascosto successivo acquisisce progressivamente maggiori dettagli man mano che i dati vengono trasmessi attraverso la rete. In un modello utilizzato per classificare immagini di animali in base alla specie, i primi livelli nascosti potrebbero semplicemente riconoscere la presenza della forma di un animale in una parte della foto; i livelli nascosti centrali potrebbero riconoscere che l'animale è un uccello; i livelli nascosti finali, appena prima del livello di output, potrebbero distinguere i dettagli specifici che differenziano una specie di uccello da un'altra specie molto simile.

L'obiettivo dei metodi di distillazione della conoscenza basati sulle caratteristiche è quindi quello di addestrare il modello studente ad apprendere le stesse caratteristiche della rete insegnante. Le funzioni di distillation loss basate sulle caratteristiche vengono utilizzate per misurare e ridurre al minimo la differenza tra le attivazioni delle caratteristiche delle due reti.

Mentre sia la conoscenza basata sulle risposte che quella basata sulle caratteristiche si concentrano sugli output di specifici livelli del modello, la distillazione della conoscenza basata sulle relazioni si concentra sulle relazioni tra diversi livelli o tra mappe di caratteristiche che rappresentano le attivazioni in diversi livelli o posizioni.

In sostanza, la conoscenza basata sulle relazioni rappresenta forse l'approccio più completo per addestrare la rete studente a emulare il "processo di pensiero" del modello insegnante. Queste relazioni e correlazioni possono essere modellate in vari modi, tra cui le correlazioni tra le mappe di caratteristiche, le matrici che rappresentano la somiglianza tra i diversi livelli, gli incorporamenti di caratteristiche o le distribuzioni probabilistiche delle rappresentazioni di caratteristiche.

Nella distillazione offline, la rete insegnante è già pre-addestrata e i pesi del modello sono congelati per impedire ulteriori modifiche. La distillazione offline è tipica di molti approcci KD per gli LLM, in cui l'insegnante è spesso un modello proprietario più grande per il quale non è possibile modificare i pesi del modello.

In alcune circostanze, potrebbe non essere disponibile un modello insegnante adeguatamente pre-addestrato e con prestazioni idonee, oppure un data scientist potrebbe voler adattare la rete insegnante al proprio caso d'uso specifico. I programmi di distillazione online mirano ad addestrare contemporaneamente sia la rete insegnante che la rete studente.

Ad esempio, Cioppa et al hanno proposto uno schema di distillazione online per i modelli di segmentazione semantica utilizzati negli eventi sportivi dal vivo, in cui le circostanze visive potrebbero cambiare durante una partita. L'obiettivo era quello di aggirare il compromesso tra la velocità di una rete più piccola e l'accuratezza di una rete più grande, addestrando in modo continuo un modello lento e ben performante sui dati delle partite in tempo reale e contemporaneamente distillando la conoscenza di quel modello più grande in un modello più piccolo e più veloce distribuito per generare output in tempo reale.⁶

Nell'autodistillazione, una rete agisce sia come insegnante che come studente. Mentre la distillazione della conoscenza convenzionale comporta il trasferimento di conoscenze da un modello all'altro, l'autodistillazione può essere intesa come il trasferimento della conoscenza dai livelli più profondi di una rete ai livelli superficiali della stessa rete.⁷

Nell'autodistillazione, più "classificatori superficiali" basati sull'attenzione vengono aggiunti ai livelli intermedi del modello a profondità variabili. Durante l'addestramento, i classificatori più profondi fungono da modelli insegnante e guidano l'addestramento degli altri moduli basati sull'attenzione attraverso due tipi di distillation loss: una KL divergence metric loss sugli output e una L2 regularization loss sulle mappe delle caratteristiche.

Una volta che il modello è stato addestrato ed è pronto per l'inferenza, tutti questi classificatori superficiali vengono rimossi dal modello. In sostanza, questo permette al modello di essere più grande e di avere una maggiore capacità di riconoscimento dei pattern durante l'addestramento, ma poi di essere più piccolo e di conseguenza più veloce ed efficiente quando viene distribuito.

Gli LLM con le funzionalità più avanzate sono, nella maggior parte dei casi, troppo costosi e richiedono troppe risorse di elaborazione per essere accessibili a molti aspiranti utenti come hobbisti, startup o istituti di ricerca. Inoltre, nonostante le loro prestazioni avanzate e le loro capacità uniche, gli LLM proprietari per loro natura non possono essere adattati ad applicazioni di nicchia e casi d'uso specifici.

Inoltre, la maggior parte degli LLM disponibili in commercio sono troppo grandi e impegnativi dal punto di vista computazionale per essere utilizzati localmente su telefoni cellulari o altri dispositivi edge. Questo presenta diverse complicazioni logistiche, computazionali e di privacy che potrebbero essere aggirate con un modello più piccolo che potrebbe essere eseguito direttamente sui dispositivi mobili. La compressione dei modelli della KD rappresenta quindi un mezzo promettente per trasferire le qualità emergenti dei modelli di grandi dimensioni a modelli sufficientemente piccoli da essere eseguiti sui dispositivi.