Toward Efficient Uncertainty in LLMs through Evidential Knowledge Distillation

Sintesi Tecnica: Verso un'incertezza efficiente negli LLM attraverso la Distillazione della Conoscenza Evidenziale

Problematica
La quantificazione accurata dell'incertezza è una sfida critica per i Large Language Models (LLM). Sebbene metodi con una solida base teorica come le Reti Neurali Bayesiane (BNN) e gli approcci basati su ensemble (ad esempio, i Bayesian Prompt Ensembles) modellino efficacemente l'incertezza epistemica e migliorino la calibrazione, essi soffrono di costi computazionali proibitivi. Questi metodi richiedono tipicamente molteplici passaggi in avanti (sampling) sia durante l'addestramento che durante l'inferenza per approssimare le distribuzioni predittive tramite campionamento Monte Carlo. Questo requisito di multi-passaggio aumenta significativamente la latenza di inferenza, limitando l'implementazione pratica di LLM consapevoli dell'incertezza in applicazioni real-time o su larga scala.

Metodologia
Gli autori propongono un framework per la Distillazione della Conoscenza Evidenziale per trasferire le stime di incertezza da modelli "insegnanti" computazionalmente costosi, basati sul sampling, a modelli "studenti" efficienti a passaggio singolo.

1. Modelli Insegnanti (Teacher Models): Lo studio utilizza due distinti insegnanti basati sul sampling:

   • Bayesian Prompt Ensembles (BayesPE): Una formulazione di ensemble su prompt di input.
   • Laplace LoRA: Un approccio bayesiano che utilizza l'approssimazione di Laplace sui pesi LoRA.
     Entrambi gli insegnanti generano distribuzioni predittive aggregando molteplici campioni (prompt o perturbazioni dei pesi) per catturare l'incertezza totale (aleatoria + epistemica).

2. Architettura dello Studente (Student Architecture): Gli studenti sono LLM pre-addestrati (Mistral 7B Instruct) sottoposti a fine-tuning tramite Low-Rank Adaptation (LoRA). Ciò consente aggiornamenti dei parametri efficienti senza riaddestrare l'intero modello. Il processo di distillazione inizializza le matrici LoRA dello studente dal docente (se applicabile) e le affina utilizzando una testa di classificazione modificata.

3. Strategie di Distillazione: Il documento confronta due parametrizzazioni di output dello studente:

   • Studente Softmax (Softmax Student): Un modello baseline che emette una singola distribuzione categorica (medie delle probabilità) tramite softmax. Minimizza la log-likelihood negativa delle predizioni medie del docente. Sebbene veloce, collassa tutta l'incertezza nell'entropia di Shannon, fallendo nel distinguere tra i componenti aleatori ed epistemici.
   • Studente Evidenziale (Dirichlet Student): Questo modello impiega il Deep Learning Evidenziale, emettendo i parametri ($\alpha$) di una distribuzione di Dirichlet. Inveve di predire direttamente le probabilità di classe, predice i parametri di concentrazione derivati dai logit ($\alpha_c = 1 + \text{softplus}(z_c)$). L'evidenza totale ($\alpha_0 = \sum \alpha_c$) funge da misura esplicita di confidenza. Questa formulazione permette al modello di separare l'incertezza aleatoria (rumore intrinseco dei dati) dall'incertezza epistemica (ignoranza del modello) all'interno di un singolo passaggio in avanti.

4. Obiettivo di Addestramento (Training Objective): Gli studenti sono addestrati per corrispondere alla distribuzione predittiva del docente. La funzione di perdita è la log-likelihood negativa delle predizioni campionate del docente sotto la distribuzione di Dirichlet predetta.

   • Regolarizzazione: Gli autori investigano varie strategie di regolarizzazione per il parametro di concentrazione $\alpha_0$, tra cui fissarlo a una costante, spingerlo verso un prior uniforme o ancorarlo a una media apprendibile dell'intero dataset, per prevenire l'overconfidence e stabilizzare l'addestramento.
   • Early Stopping: L'addestramento viene interrotto in base alla log-likelihood negativa (NLL) sulle etichette reali per evitare l'overfitting sul rumore del docente.

Risultati Chiave
Gli esperimenti sono stati condotti su quattro dataset di classificazione del testo (Amazon, SST2, Yahoo, YouTube) utilizzando il backbone Mistral 7B.

• Efficienza di Inferenza: Gli studenti distillati ottengono una riduzione massiccia del tempo di inferenza rispetto ai loro docenti. Per i docenti BayesPE, lo studente Dirichlet è da 11 a 36 volte più veloce (ad esempio, un'accelerazione di 36x su Yahoo), poiché richiede un solo passaggio in avanti rispetto ai molteplici passaggi richiesti dal docente.
• Performance Predittiva: Gli studenti distillati eguagliano o occasionalmente superano l'accuratezza predittiva dei loro docenti. Ad esempio, lo studente Dirichlet ha superato il docente BayesPE sui dataset Yahoo (+1.7 pp) e YouTube (+2.5 pp).
• Quantificazione dell'Incertezza:
  • Calibrazione: Gli studenti Dirichlet ottengono generalmente l'Expected Calibration Error (ECE) e i punteggi Brier più bassi rispetto sia ai docenti che alle baseline Softmax.
  • Rilevamento OOD (Out-of-Distribution): Negli esperimenti OOD, lo studente Dirichlet ha dimostrato una capacità superiore di distinguere tra campioni in-domain e OOD (AUROC più elevata) rispetto al docente e allo studente Softmax. Ha avuto successo nell'amplificare l'entropia predittiva totale per i dati OOD.
  • Decomposizione: A differenza dello studente Softmax, lo studente Dirichlet riesce a scomporre l'incertezza in componenti aleatorie ed epistemiche, fornendo segnali di incertezza più ricchi.
• Convergenza: Gli studenti distillati da docenti BayesPE hanno raggiunto la convergenza molto rapidamente (spesso entro 1–5 epoche), mentre quelli da docenti Laplace hanno richiesto più epoche (20–40), suggerendo un panorama di ottimizzazione più complesso per questi ultimi.

Significatività e Claim
Il documento afferma di essere la prima applicazione del Deep Learning Evidenziale alla distillazione di Large Language Models. Il suo contributo primario è l'instaurazione di una base per la stima dell'incertezza efficiente a passaggio singolo negli LLM.

• Trade-off tra Efficienza e Accuratezza: Il lavoro dimostra che è possibile ereditare la robusta modellazione dell'incertezza dei metodi di sampling computazionalmente intensivi senza incorrere nella relativa latenza di inferenza.
• Modellazione Epistemica: Utilizzando gli output di Dirichlet, il metodo modella esplicitamente l'incertezza epistemica, una capacità persa nella standard distillazione softmax.
• Ambito e Limitazioni: Gli autori inquadrano modestamente il presente lavoro come una prova di concetto limitata ai compiti di classificazione di testo discreto. Riconoscono che, sebbene i risultati siano promettenti, estendere questo framework alla generazione di testo aperto rimane una direzione per il lavoro futuro. Lo studio si concentra sul far corrispondere le performance del docente piuttosto che sul migliorare le capacità assolute dei docenti, sebbene i risultati empirici mostrino che gli studenti a volte superano i docenti grazie agli effetti regolarizzanti della distillazione e dell'early stopping.