CLSGen: A Dual-Head Fine-Tuning Framework for Joint Probabilistic Classification and Verbalized Explanation

1. Il Problema

L'articolo affronta una sfida critica nell'applicazione dei Large Language Models (LLM) a compiti di decisione ad alto rischio, come la previsione degli esiti clinici o la scoperta di farmaci. Sebbene gli LLM moderni siano eccellenti nel generare spiegazioni verbali e nel gestire contesti lunghi, presentano due limiti fondamentali per l'uso pratico:

1. Mancanza di stime probabilistiche affidabili: Gli LLM faticano a fornire stime quantitative di probabilità (confidenza) robuste.
2. Il paradosso del "Catastrophic Forgetting" (Dimenticanza Catastrofica): I tentativi di adattare gli LLM (decoder-only) a compiti di classificazione binaria tramite tecniche tradizionali (aggiunta di un layer Sigmoid e ottimizzazione con Binary Cross-Entropy) portano spesso al collasso linguistico. In questo scenario, il modello perde le sue capacità generative originali, producendo testo incoerente o privo di senso, rendendo di fatto inutile la sua capacità di fornire spiegazioni verbali.

2. Metodologia: CLSGen

Per superare questi limiti, gli autori propongono CLSGen, un nuovo framework di fine-tuning progettato per compiti di classificazione binaria che mantiene simultaneamente la capacità di stima probabilistica e quella di generazione di spiegazioni.

L'approccio si articola in tre componenti principali:

A. Architettura del Modello (Dual-Head)

CLSGen utilizza un modello di base decoder-only (es. Llama-3.1-8B) equipaggiato con una struttura a doppia testa:

• Testa di Generazione (LM Head): Un layer standard di language modeling che genera la sequenza di testo, inclusa la spiegazione (ragionamento) e l'etichetta verbalizzata (es. "Classification: 1:death").
• Testa di Classificazione (CLS Head): Un layer aggiuntivo (un MLP a due strati seguito da un layer Sigmoid) che estrae la rappresentazione nascosta dell'ultimo token dell'input e predice direttamente l'etichetta binaria (0 o 1).

Il modello viene addestrato minimizzando una funzione di perdita totale composta dalla somma di due termini:
$$\mathcal{L}_{total} = \mathcal{L}_{gen} + \mathcal{L}_{cls}$$
Dove $\mathcal{L}_{gen}$ è la perdita di modellazione causale (per la generazione del testo) e $\mathcal{L}_{cls}$ è la perdita di classificazione. Questa ottimizzazione congiunta, effettuata tramite LoRA (Low-Rank Adaptation), permette ai gradienti di entrambi gli obiettivi di plasmare le rappresentazioni condivise senza distruggere le capacità linguistiche.

B. Costruzione dei Dati

Per garantire la qualità delle spiegazioni e la loro allineamento con le etichette vere (ground-truth), gli autori adottano una strategia di apprendimento semi-supervisionato:

1. Un LLM esterno ($M_{data}$) genera spiegazioni e previsioni per ogni istanza di input.
2. Vengono eseguiti più tentativi di generazione ($K=5$).
3. Vengono conservati solo i campioni in cui la spiegazione generata è valida e la previsione verbale corrisponde all'etichetta vera.
4. Il dataset finale è un insieme aumentato $T = \{(x_i, y_i, e_i)\}$, dove $e_i$ è la spiegazione giustificativa.

C. Addestramento

Il modello viene fine-tuned su questo dataset aumentato. Durante l'inferenza, il modello riceve un prompt strutturato e produce sia una spiegazione testuale che una probabilità numerica (tramite la testa di classificazione), garantendo coerenza tra i due output.

3. Contributi Chiave

• Framework CLSGen: Una nuova architettura che risolve il compromesso tra accuratezza nella classificazione e capacità di generazione, evitando il collasso linguistico tipico del fine-tuning discriminatorio puro.
• Stime Probabilistiche e Spiegazioni: Il modello fornisce non solo una classificazione, ma anche una stima di probabilità affidabile (AUROC) e una giustificazione testuale coerente con tale probabilità.
• Validazione Empirica del Collasso Linguistico: Gli autori dimostrano sperimentalmente che il fine-tuning per sola classificazione porta alla perdita della capacità di seguire le istruzioni (misurata tramite il tasso di output "parsable"), giustificando la necessità di un approccio ibrido.
• Metodologia di Valutazione Ibrida: Introduzione di nuove metriche per valutare la coerenza tra la testa di classificazione e il testo generato, come AUROC-Alignment e Cohen's Kappa tra etichette numeriche e verbali.

4. Risultati Sperimentali

Il framework è stato valutato su benchmark reali e complessi:

• Dataset: LCD Benchmark (previsione di mortalità a 30/90 giorni da note di dimissione) e CliniFact (verifica di affermazioni cliniche).
• Metriche di Classificazione: CLSGen supera i metodi basati solo sull'inferenza (prompting verbale, auto-consistenza) e i modelli CNN tradizionali. In particolare, ottiene i migliori punteggi AUROC (es. 0.8586 su LCD 30 giorni vs 0.7455 per l'auto-consistenza) e F1-score.
• Coerenza e Leggibilità:
  • AUROC-Alignment: Punteggi molto alti (0.95 - 0.98), indicando che la probabilità calcolata dalla testa di classificazione è fortemente allineata con l'etichetta verbalizzata nel testo.
  • Coerenza Logica: L'uso di un "LLM-as-a-judge" ha mostrato che oltre il 99.8% delle spiegazioni generate su LCD è logicamente coerente con l'etichetta predetta (RLI < 0.002).
  • Parsability: Il modello mantiene un tasso di output parsabili vicino al 100%, dimostrando che la struttura del testo non è andata persa.

5. Significato e Impatto

Il lavoro di CLSGen è significativo perché offre una soluzione pratica per l'adozione degli LLM in ambiti critici come la sanità.

• Affidabilità Decisionale: Permette di ottenere stime di probabilità quantitative (essenziali per la medicina basata sull'evidenza) senza sacrificare la trasparenza offerta dalle spiegazioni testuali.
• Stabilità: Risolve il problema della "linguistic collapse", garantendo che il modello rimanga un generatore di testo utile anche dopo essere stato specializzato per compiti discriminativi.
• Interpretabilità: La forte allineamento tra la probabilità numerica e la giustificazione testuale aumenta la fiducia degli utenti umani nelle decisioni automatizzate, un requisito fondamentale per l'implementazione clinica reale.

In sintesi, CLSGen dimostra che è possibile addestrare modelli LLM per compiti di classificazione ad alta precisione mantenendo intatte le loro capacità generative, creando sistemi di IA più sicuri, interpretabili e pronti per l'uso nel mondo reale.