Reinforcement Learning Improves LLM Accuracy and Reasoning in Disease Classification from Radiology Reports

1. Il Problema

La classificazione accurata delle malattie dai referti radiologici è fondamentale per molte applicazioni cliniche. Sebbene il Fine-Tuning Supervisionato (SFT) di modelli linguistici leggeri (Lightweight LLMs) abbia dimostrato di migliorare l'accuratezza nella classificazione, introduce un limite critico: degrada la capacità di ragionamento esplicito.
Nel contesto medico, la trasparenza e l'esplicabilità sono cruciali per la fiducia clinica e il processo decisionale. Lo SFT tradizionale, che si basa solo su etichette di malattia (senza supervisione sul processo di ragionamento), tende a far sì che i modelli producano liste di diagnosi senza fornire le giustificazioni o le prove tratte dal referto, rendendo le raccomandazioni algoritmiche poco affidabili per i medici. Inoltre, l'uso di API esterne solleva preoccupazioni sulla privacy, mentre l'hosting locale di modelli grandi richiede risorse computazionali eccessive, rendendo necessari modelli leggeri ottimizzati.

2. Metodologia Proposta

Gli autori propongono un framework di training e inferenza in due fasi progettato per migliorare sia l'accuratezza che la capacità di ragionamento dei modelli leggeri, senza richiedere dati di supervisione espliciti per il ragionamento nella fase finale.

Fase 1: Supervised Fine-Tuning (SFT)

• Obiettivo: Adattare il modello al dominio radiologico.
• Dati: Utilizzo di 2.000 referti dal dataset MIMIC-CXR con etichette di malattia generate da GPT-4o.
• Processo: Il modello viene addestrato solo sulle etichette di malattia (output: lista di patologie). La sezione di ragionamento nel prompt di output viene lasciata intenzionalmente vuota durante questo stadio, poiché non c'è supervisione esplicita sul ragionamento.

Fase 2: Ottimizzazione della Politica Relativa di Gruppo (GRPO)

• Obiettivo: Raffinare le previsioni ottimizzando accuratezza e formato, ripristinando e potenziando il ragionamento.
• Algoritmo: Viene applicato GRPO (un algoritmo di Reinforcement Learning) sullo stesso dataset della Fase 1.
• Funzione di Ricompensa (Reward Function): A differenza degli approcci precedenti, non richiede etichette di ragionamento umane. La ricompensa è basata su regole:
  1. Accuratezza di classificazione: Calcolata come media di precisione e recall rispetto alle etichette di verità.
  2. Correttezza del formato: Assegna un reward zero se la sezione di ragionamento è vuota o se ripete semplicemente la lista delle malattie. Questo forza il modello a generare un ragionamento esplicito e coerente basato su frasi del referto.
  • Formula: $Reward = 0.8 \times Accuratezza + 0.2 \times Formato$.

Inferenza e Post-Processing

• Inferenza Multipla (Ensemble): Per ogni referto, il modello genera 5 inferenze indipendenti.
• Voto di Maggioranza: Le etichette finali delle malattie sono determinate tramite voto di maggioranza tra le 5 inferenze, aumentando la robustezza.
• Sintesi del Ragionamento: Un modello LLM leggero originale (non fine-tuned) viene utilizzato per sintetizzare i 5 ragionamenti generati in un unico rapporto coerente, giustificando la lista finale di malattie.

3. Contributi Chiave

1. Framework RL senza supervisione esplicita sul ragionamento: Dimostrano che è possibile migliorare le capacità di ragionamento di un LLM utilizzando solo una funzione di ricompensa basata su regole (accuratezza + formato), evitando la costosa annotazione umana di catene di pensiero (Chain-of-Thought).
2. Mitigazione dell'oblio catastrofico: Lo SFT tende a cancellare le capacità di ragionamento preesistenti (specialmente in modelli come LLaMA). L'aggiunta di GRPO ripristina e potenzia queste capacità.
3. Strategia di Inferenza Robusta: L'uso combinato di inferenze multiple, voto di maggioranza e sintesi del ragionamento migliora significativamente sia l'accuratezza della classificazione che la qualità della spiegazione.
4. Validazione su Modelli Leggeri: Il metodo è stato testato con successo su modelli di piccole dimensioni (LLaMA 3.1-8B, Qwen 2.5-3B, Phi-3 Mini-128K), rendendoli adatti per il deployment on-premise.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su tre dataset annotati da radiologi: MIMIC-CXR, NIH-CXR e MIDRC.

• Accuratezza di Classificazione:

  • L'SFT ha superato significativamente i modelli baseline.
  • L'aggiunta di GRPO ha portato ulteriori miglioramenti nell'F1-score micro-averaged in 8 dei 9 cohort testati.
  • Esempi di miglioramento: Qwen su MIMIC-CXR (+13.2% F1), Phi3 su NIH-CXR (+12.8% F1), LLaMA su MIDRC (+8.0% F1).
  • LLaMA con SFT+GRPO ha raggiunto un'accuratezza quasi pari a quella di GPT-4o (usato come modello "insegnante" per le etichette).

• Qualità del Ragionamento:

  • Sono stati introdotti due metriche: Reasoning Recall (quante malattie reali sono menzionate nel ragionamento) e Reasoning Comprehensiveness (quante malattie previste sono supportate da ragionamenti espliciti).
  • Risultato Critico: Dopo lo SFT, modelli come LLaMA hanno perso completamente la capacità di ragionamento (Recall = 0). L'applicazione di GRPO ha ripristinato parzialmente o totalmente questa capacità, con miglioramenti significativi in 8 dei 9 cohort.
  • Il modello Qwen ha mostrato i guadagni relativi più grandi nel ragionamento.

• Confronto con altri approcci:

  • Il metodo proposto (SFT + GRPO) ha superato o è stato comparabile all'addestramento SFT su un piccolo subset di dati con ragionamento esplicito, ma senza richiedere tale supervisione esplicita nella fase finale.
  • I modelli basati su BERT (BioClinicalBERT) hanno mostrato prestazioni inferiori nella capacità di generare ragionamenti naturali, sebbene avessero buone metriche di classificazione.

5. Significato e Implicazioni

Questo studio dimostra che l'apprendimento per rinforzo (in particolare GRPO) può colmare il divario tra l'alta accuratezza della classificazione e la necessità di spiegabilità nei modelli medici leggeri.

• Impatto Clinico: Fornisce uno strumento che non solo "indovina" la diagnosi, ma spiega il perché basandosi sulle evidenze del referto, aumentando la fiducia dei radiologi.
• Efficienza: Permette di utilizzare modelli leggeri (che possono girare su hardware locale) con prestazioni paragonabili a modelli molto più grandi o a GPT-4o, risolvendo problemi di privacy e costi computazionali.
• Futuro: Il lavoro suggerisce che la supervisione esplicita del ragionamento non è sempre necessaria se si utilizzano funzioni di ricompensa intelligenti, aprendo la strada a metodi scalabili per l'addestramento di AI medica spiegabile.

Limitazioni menzionate: La funzione di reward non supervisiona esplicitamente la lunghezza o l'allineamento semantico profondo del ragionamento; la valutazione del ragionamento è ancora affidata in parte ad altri LLM (GPT-4o/Gemini) che potrebbero introdurre bias; e la stima dell'incertezza delle previsioni rimane una sfida aperta.