Functional-level Uncertainty Quantification for Calibrated Fine-tuning on LLMs

Riepilogo Tecnico: Quantificazione dell'Incertezza a Livello Funzionale per il Fine-Tuning Calibrato su LLM

Enunciato del Problema
Una quantificazione accurata dell'incertezza (UQ) è fondamentale per il dispiegamento affidabile dei Modelli Linguistici su Larga Scala (LLM), in particolare in domini critici per la sicurezza come l'assistenza sanitaria e il ragionamento scientifico. Tuttavia, gli LLM fine-tunati mostrano frequentemente un'eccessiva sicurezza quando adattati a compiti con dati limitati, non riuscendo a riflettere la vera credibilità delle loro uscite. I metodi UQ esistenti per il Fine-Tuning Efficiente in Parametri (PEFT), come quelli basati sull'Adattamento a Basso Rango (LoRA), sono in gran parte post hoc. Stimano l'incertezza dopo che il modello è stato addestrato, analizzando lo spazio dei parametri (ad esempio tramite ensemble o approssimazioni bayesiane) o perturbando i prompt. Questi approcci spesso non affrontano la causa radice dell'eccessiva sicurezza: l'incapacità del modello di catturare l'incertezza indotta dal fine-tuning e i gap di generalizzazione durante lo stesso processo di adattamento. Inoltre, i metodi che si basano su molteplici passaggi in avanti o su campionamenti estesi introducono un sovraccarico computazionale significativo, limitando la scalabilità.

Metodologia: UQ4CT
Gli autori propongono la Quantificazione dell'Incertezza a Livello Funzionale per il Fine-Tuning Calibrato (UQ4CT), un framework che sposta il focus dall'incertezza nello spazio dei parametri all'incertezza nello spazio funzionale. Invece di analizzare i parametri appresi dopo l'addestramento, UQ4CT calibra l'incertezza sullo spazio funzionale indotto da miscele di esperti LoRA dipendenti dal prompt durante il processo di fine-tuning.

La metodologia è composta da tre componenti fondamentali:

1. Decomposizione Funzionale Gerarchica tramite MixLoRA:
   Per evitare l'esplosione combinatoria della modellazione di tutte le possibili combinazioni di esperti attraverso i livelli, UQ4CT impiega una decomposizione gerarchica. Utilizza un'architettura Mixture-of-Experts (MoE) in cui, a ogni livello, un router sparsamente dinamico seleziona un sottoinsieme di esperti LoRA basato sullo stato nascosto di input. Ciò costruisce uno spazio funzionale flessibile in cui l'output del modello è una somma pesata di diverse funzioni di base (esperti LoRA) condizionate all'input. Lo spazio funzionale è rappresentato come una composizione ricorsiva di miscele livello per livello, permettendo al modello di comporre adattivamente i percorsi di calcolo.

2. Stima della Confidenza a Livello Funzionale (FLC):
   Il modello quantifica l'incertezza analizzando la concentrazione della selezione degli esperti. Nello specifico, la Confidenza a Livello Funzionale (FLC) è derivata dai pesi di routing dei top-$k$ esperti selezionati dal router a ogni livello. Un'alta concentrazione nei pesi di routing (cioè, il router è sicuro nella selezione di esperti specifici) corrisponde a un'alta confidenza funzionale, mentre pesi dispersi indicano un'incertezza maggiore. Ciò fornisce un segnale interno e principiale della certezza del modello senza richiedere molteplici passaggi di campionamento.

3. Funzione di Perdita di Calibrazione ($L_{cal}$):
   L'innovazione centrale è una nuova funzione di perdita di calibrazione che allinea la confidenza a livello funzionale con la correttezza predittiva durante l'addestramento. La perdita è definita come:
   $$L_{cal} = (1\{\text{MixLoRA}(x) = y^*\} - \text{FLC}(x))^2$$
   dove $1\{\cdot\}$ è una funzione indicatrice per le previsioni corrette. Questo obiettivo forza la FLC ad approssimare la probabilità che la previsione del modello sia corretta. Minimizzando questa perdita insieme alle perdite standard di entropia incrociata e di bilanciamento del carico, il modello impara ad assegnare alta confidenza solo quando è probabile che sia corretto e un'incertezza maggiore quando non lo è. Ciò incoraggia il router a esplorare miscele di esperti alternative quando l'incertezza è alta, migliorando così la generalizzazione.

Contributi Chiave

• Prospettiva a Livello Funzionale: Il documento introduce un passaggio dalla quantificazione dell'incertezza nello spazio dei parametri a quello nello spazio funzionale, calibrando esplicitamente la confidenza del modello sulle mappature input-output che può realizzare durante il fine-tuning.
• Fine-Tuning Consapevole della Calibrazione: Viene proposta una nuova funzione di perdita di calibrazione che allinea dinamicamente i pesi delle miscele LoRA dipendenti dal prompt con la correttezza predittiva, mitigando l'eccessiva sicurezza senza degradare l'accuratezza.
• Architettura Efficiente: Il metodo sfrutta un'architettura gerarchica MixLoRA che cattura relazioni funzionali diverse con un sovraccarico computazionale minimo, richiedendo un solo passaggio in avanti durante l'inferenza.
• Garanzia Teorica: Gli autori forniscono una dimostrazione che l'ottimizzazione del rischio di calibrazione sulla distribuzione dei dati produce una FLC ottimale che corrisponde alla vera probabilità di correttezza predittiva.

Risultati Sperimentali
Gli autori hanno valutato UQ4CT su quattro benchmark a scelta multipla (ARC-E, ARC-C, OBQA, ClimateQA) e due compiti di QA generativa a risposta aperta (TruthfulQA, TriviaQA) utilizzando Llama-3.1-8B e Mistral-7B.

• Prestazioni In-Distribution: UQ4CT ha ridotto l'Errore di Calibrazione Atteso (ECE) di oltre il 25% su tutti i benchmark rispetto a baseline solide (inclusi LoRA, MC Dropout, Deep Ensembles e Bayesian LoRA), mantenendo un'accuratezza competitiva o superiore. Notevolmente, UQ4CT ha ottenuto questi risultati con un singolo passaggio in avanti, mentre metodi come BLoB (N=10) richiedevano dieci passaggi e tuttavia producevano un ECE più alto.
• Shift di Distribuzione: Sotto shift di distribuzione (testati su sottocompiti MMLU in Informatica, Ingegneria, Diritto e Salute), UQ4CT ha mantenuto una calibrazione superiore e un'accuratezza competitiva, dimostrando una maggiore robustezza agli input fuori distribuzione rispetto alle baseline.
• Compiti Generativi: Nella generazione a risposta aperta, UQ4CT ha ottenuto il più basso ECE su TruthfulQA e TriviaQA, suggerendo che la prospettiva di calibrazione nello spazio funzionale si estende oltre gli scenari di tipo classificazione.
• Efficienza: L'analisi della memoria indica che l'uso di picco della memoria GPU di UQ4CT è solo ~7,5% superiore rispetto al LoRA standard, rendendolo computazionalmente fattibile per il dispiegamento su larga scala.

Significato e Affermazioni
Il documento afferma che UQ4CT affronta una limitazione fondamentale nelle attuali pratiche di fine-tuning: la mancanza di consapevolezza dell'incertezza durante il processo di adattamento. Integrando la calibrazione dell'incertezza direttamente nell'obiettivo di fine-tuning tramite segnali a livello funzionale, il metodo garantisce che la confidenza del modello rifletta meglio la sua reale affidabilità predittiva. Gli autori posizionano questo lavoro come un passo fondamentale verso LLM fine-tunati affidabili, in particolare in ambienti con scarsità di dati dove l'eccessiva sicurezza comporta rischi significativi. Notano con modestia che, sebbene l'approccio si estenda alla generazione a risposta aperta utilizzando proxy basati su F1, la scalabilità al ragionamento a lungo raggio o alla generazione multi-step rimane una direzione per lavori futuri.