Domain-Shift-Aware Conformal Prediction for Large Language Models

Sintesi Tecnica: Conformal Prediction Sensibile al Domain-Shift per i Large Language Models

Definizione del Problema
I Large Language Models (LLM) producono frequentemente output eccessivamente sicuri di sé ma fattualmente errati (allucinazioni), ponendo rischi significativi in domini ad alto rischio come la sanità e la finanza. Sebbene la Conformal Prediction (CP) offra un framework principato per la quantificazione dell'incertezza (UQ) con garanzie di copertura a campioni finiti e indipendenti dalla distribuzione, la CP standard si basa sull'assunzione di scambiabilità. Tale assunzione richiede che i dati di calibrazione e di test siano estratti dalla stessa distribuzione. Nelle implementazioni reali, il domain shift (variazione di dominio) è pervasivo; i prompt di test differiscono spesso significativamente dai dati di calibrazione storici a causa dell'evoluzione del comportamento degli utenti o delle distribuzioni di input. Sotto tali shift, la CP standard soffre frequentemente di sottocopertura (under-coverage), producendo set di predizione troppo stretti per contenere la risposta corretta, fallendo così nel fornire affidabili garanzie di incertezza.

Metodologia: Framework DS-CP
Gli autori propongono la Domain-Shift-Aware Conformal Prediction (DS-CP), un framework progettato per adattare la CP agli LLM sotto domain shift senza richiedere dati etichettati nel target. Il metodo affronta l'infeasibilità di stimare i rapporti di densità in spazi di prompt ad alta dimensionalità e non strutturati attraverso i seguenti passaggi:

1. Embedding Semantico: Invece di operare nello spazio dei token grezzi, il framework utilizza un modello di embedding pre-addestrato $h: \mathcal{X} \to \mathcal{Z}$ per proiettare i prompt in uno spazio semantico a dimensionalità ridotta. Ciò cattura la similarità cross-dominio, permettendo al domain shift di essere caratterizzato dallo shift nella distribuzione degli embedding.
2. Stima del Rapporto di Densità: Nello spazio degli embedding, viene addestrato un classificatore di dominio per distinguere tra i prompt sorgente (calibrazione) e quelli target (test). Questo classificatore stima il rapporto di densità $\tilde{r}(z) = \frac{dP'_Z}{dP_Z}(z)$, che funge da peso per i campioni di calibrazione in base alla loro vicinanza al dominio target.
3. Pesatura Regolarizzata: Per evitare che i set di predizione degenerino nell'intero spazio di output (un rischio quando i rapporti di densità sono altamente sbilanciati a causa di grandi shift o piccole dimensioni del campione), gli autori introducono un passaggio di regolarizzazione. Il peso assegnato al punto di test è sostituito da un parametro regolarizzato $\lambda$ che dipende solo dai dati di calibrazione.
   • La distribuzione del punteggio empirico è costruita come:
     $$\hat{\mu}_n = \sum_{i=1}^n \frac{\hat{w}_i}{\sum \hat{w}_j + \lambda} \delta_{S_i} + \frac{\lambda}{\sum \hat{w}_j + \lambda} \delta_{\infty}$$
     dove $\hat{w}_i$ sono i rapporti di densità stimati per i campioni di calibrazione, e $\lambda$ è scelto (tipicamente come $\max_i \hat{w}_i$) per bilanciare validità e adattività.
4. Costruzione del Set di Predizione: Il set di predizione finale per un nuovo prompt è definito dal quantile $(1-\alpha)$ di questa distribuzione empirica pesata.

Garanzie Teoriche
Il paper stabilisce limiti di copertura a campioni finiti per DS-CP:

• Limite Inferiore: La probabilità di copertura è garantita essere almeno $1 - \alpha$ meno un termine di scarto. Questo scarto dipende dai pesi normalizzati e dalla distanza di variazione totale tra le distribuzioni dei punteggi dei punti di calibrazione e di test nello spazio degli embedding.
• Condizione di Smoothness: Gli autori dimostrano che se il punteggio di non conformità varia in modo fluido rispetto all'embedding semantico, il gap di copertura rimane piccolo anche se le distribuzioni dei prompt grezzi sono distinte.
• Conservatività: Viene fornito un limite superiore, mostrando che il termine di regolarizzazione $\lambda$ introduce un grado controllato di conservatività (sovra-copertura) quando la dimensione effettiva del campione dei dati di calibrazione rilevanti è bassa.
• Recupero della CP Standard: Sotto la condizione di scambiabilità (assenza di domain shift), il metodo si riduce esattamente alla CP standard, garantendo l'assenza di perdita di validità nell'impostazione ideale.

Risultati Sperimentali
Il framework è stato valutato sul benchmark MMLU, trattando diverse categorie di materie come domini distinti per simulare il domain shift. Gli esperimenti sono stati condotti su 16 LLM open-source appartenenti a 9 famiglie di modelli.

• Copertura: La CP standard ha spesso fallito nel raggiungere la copertura target del 90%, scendendo spesso significativamente al di sotto di questo livello quando i domini di calibrazione e test divergevano. DS-CP ha migliorato costantemente la copertura, mantenendola più vicina al livello target in tutti i modelli e coppie di domini.
• Dimensione del Set: Sebbene DS-CP abbia prodotto set di predizione leggermente più ampi rispetto alla CP standard per compensare lo shift, l'incremento è stato modesto. Ciò dimostra un riuscito compromesso tra il ripristino della validità e il mantenimento dell'efficienza.
• Adattività: L'analisi ha mostrato che DS-CP corregge selettivamente la sottocopertura nei casi difficili (dove la CP standard fallisce) pur apportando aggiustamenti minimi quando la CP standard è già valida, provando che non si tratta semplicemente di un ampliamento uniforme dei set.

Significatività e Rivendicazioni
Il paper sostiene di fornire un passo pratico verso una quantificazione dell'incertezza affidabile per gli LLM in fase di deployment reale. La sua principale significatività risiede in:

1. Estensione della CP oltre la scambiabilità: Offre un metodo sistematico e teoricamente fondato per gestire il domain shift negli LLM, uno scenario in cui la CP standard fallisce.
2. Adattamento Semantico: Sfruttando gli embedding semantici invece del testo grezzo, rende la stima del rapporto di densità fattibile per gli input degli LLM ad alta dimensionalità.
3. Affidabilità Bilanciata: Raggiunge un equilibrio tra validità statistica (preservando le garanzie di copertura) e adattività (personalizzando i set di predizione al dominio target specifico), affrontando la questione critica delle predizioni inaffidabili nelle applicazioni critiche per la sicurezza.

Gli autori segnalano dei limiti, tra cui la dipendenza dalla qualità degli estimator di embedding e di rapporto di densità, la necessità di prompt target non etichettati e l'attuale focalizzazione su task a scelta multipla piuttosto che sulla generazione open-ended. Tuttavia, essi sostengono che DS-CP rappresenti un'adattazione fondata della conformal prediction specificamente progettata per la natura ad alta dimensionalità e non strutturata dei prompt degli LLM sotto distribuzione shift.