Riepilogo Tecnico: Modelli di Linguaggio di Grandi Dimensioni con Augmentation di Knowledge Graph per la Predizione di Malattie

Definizione del Probleamento

Le cartelle cliniche elettroniche (EHR) contengono dati longitudinali adatti alla predizione clinica, tuttavia i modelli esistenti spesso non riescono a fornire un ragionamento a livello di paziente che sia azionabile per i clinici. I modelli classici (ad es. regressione logistica, modelli ensemble ad albero) offrono robustezza e interpretabilità globale, ma tipicamente forniscono solo spiegazioni post-hoc grossolane (ad es. punteggi di importanza delle caratteristiche) che mancano della narrazione passo dopo passo necessaria per specifiche decisioni cliniche. Al contrario, sebbene i Large Language Models (LLM) possano generare narrazioni di tipo Chain-of-Thought (CoT), la loro applicazione non vincolata in medicina è ostacolata dalla scarsità di dataset di ragionamento di alta qualità e validati clinicamente, nonché dal rischio di passaggi intermedi allucinati o medicalmente errati. La sfida centrale è consentire agli LLM di eseguire una predizione di malattie accurata a livello di visita, generando al contempo spiegazioni che siano fattualmente fondate sulla conoscenza biomedica e temporalmente coerenti con la storia del paziente.

Metodologia

Gli autori propongono un framework KG-guided CoT progettato per strutturare il ragionamento dell'LLM utilizzando la struttura di PrimeKG, un knowledge graph biomedico su larga scala. La pipeline consiste in quattro fasi primarie:

1. Mappatura delle Entità del KG:

   • I codici ICD-9 da MIMIC-III sono allineati ai nodi di PrimeKG tramite una procedura in tre stadi:
     • Corrispondenza Esatta (Exact Match): Corrispondenza diretta tra le descrizioni ICD-9 e le etichette dei nodi del KG.
     • Corrispondenza per Similarità (Similarity Match): Recupero basato su embedding (similarità del coseno > 0.85) per corrispondenze non esatte.
     • Filtraggio tramite LLM: GPT-4o valida o revisiona le mappature provvisorie per garantire la correttezza clinica.
   • Questo processo produce 1.513 nodi di caratteristiche mappati e 10 nodi di malattie mappati.

2. Ricerca della Rilevanza e dei Percorsi nel KG:

   • Per ogni malattia target, il sistema identifica un insieme compatto di nodi di caratteristiche rilevanti ($R_d$) utilizzando GPT-4o.
   • Percorsi di Ragionamento (Reasoning Paths): Vengono estratti i percorsi più brevi (limitati a $L=5$ hop) tra i nodi di caratteristiche rilevanti e il nodo della malattia target all'interno di PrimeKG.
   • Selezione dei Percorsi: GPT-4o filtra questi percorsi per trattenere solo quelli clinicamente rilevanti e meccanicisticamente solidi, eliminando le catene spurie.

3. Generazione e Filtraggio della CoT:

   • L'LLM viene istruito per generare spiegazioni CoT ($C_{t,d}$) per una specifica visita $t$, incorporando le caratteristiche ICD-9 osservate, l'insieme di rilevanza specifica della malattia e i percorsi del KG estratti.
   • Filtraggio della Coerenza del Label (Label-Consistency Filtering): Le tracce CoT generate vengono trattenute solo se la conclusione binaria esplicita ($\hat{y}$) derivata dal ragionamento coincide con l'esito reale ($y$). Ciò crea un corpus di supervisione di alta qualità e fondato sul KG.

4. Fine-Tuning del Modello:

   • Modelli leggeri a pesi aperti (LLaMA-3.1-Instruct-8B e Gemma-7B) vengono sottoposti a fine-tuning su questo corpus curato.
   • In fase di inferenza, questi modelli prendono in input le caratteristiche ICD-9 e l'evidenza specifica della malattia dal KG per produrre sia una predizione della malattia che una traccia di ragionamento in stile clinico.

Contributi Chiave

• Framework KG-Guided CoT: Un metodo innovativo che integra la struttura del KG biomedico direttamente nel processo di generazione della CoT, andando oltre le spiegazioni a forma libera non vincolata verso un ragionamento clinicamente fondato.
• Supervisione Efficiente nei Dati: La costruzione di un corpus di supervisione compatto e di alta qualità tramite il filtraggio della coerenza del label permette a piccoli LLM (8 miliardi di parametri) di apprendere il ragionamento prognostico con dati limitati (400–1.000 campioni).
• Trasferimento Zero-Shot: Il framework dimostra la capacità di trasferire i pattern di ragionamento appresi a una coorte distinta (CRADLE) senza ri-addestramento, mantenendo un'alta accuratezza nella previsione di eventi cardiovascolari post-diagnosi di diabete.
• Valutazione Centrata sul Clinico: Uno studio umano in cieco valida che le spiegazioni guidate dal KG sono preferite rispetto alle baseline in termini di chiarezza, rilevanza e correttezza clinica.

Risultati Sperimentali

Il framework è stato valutato sul cohort MIMIC-III (10 malattie mappate su PrimeKG) e sul cohort CRADLE (diabete di tipo 2 verso eventi cardiovascolari).

• Performance su MIMIC-III:

  • In scenari con pochi dati (400 e 1.000 visite indice), il fine-tuning KG-guided CoT ha migliorato le prestazioni di LLaMA-3.1-8B e Gemma-7B, ottenendo risultati competitivi o superiori alle forti baseline classiche (XGBoost, Random Forest, Regressione Logistica).
  • LLaMA-3.1 (1.000 campioni): Ha raggiunto un AUROC di 0.6995 e un macro-AUPR di 0.4685, superando leggermente la migliore baseline classica (Regressione Logistica: AUROC 0.6874, AUPR 0.4178; XGBoost: AUROC 0.6799).
  • Gemma-7B: Ha mostrato miglioramenti simili, con il macro-AUPR che è aumentato da 0.25 (originale) a 0.4363 (dopo il tuning KG-CoT), sebbene l'AUROC (0.6609) sia rimasto leggermente al di sotto delle migliori baseline classiche.
  • Ablazione: La rimozione della guida del KG o del filtraggio ha ridotto le prestazioni, confermando che sia il fondamento strutturale che il filtro di coerenza del label sono critici.

• Trasferimento Zero-Shot a CRADLE:

  • I modelli sottoposti a fine-tuning su MIMIC-III sono stati applicati a CRADLE senza ri-addestramento.
  • L'accuratezza è migliorata da 0.40–0.51 (non tunato) a 0.72–0.77 (tunato con KG-CoT).
  • Sebbene l'AUROC e l'AUPR abbiano visto guadagni modesti, il significativo salto nell'accuratezza suggerisce che i modelli abbiano appreso un confine decisionale più robusto che generalizza attraverso le coorti.

• Valutazione Umana:

  • In uno studio in cieco con 115 casi, i clinici hanno preferito il modello guidato dal KG rispetto alla baseline non tunata nel 96.5% dei casi per Chiarezza/Coerenza, nel 94.8% per Copertura/Rilevanza e nel 98.3% per Correttezza/Solidità.
  • I clinici hanno notato che le tracce guidate dal KG erano meglio organizzate temporalmente, evidenziavano fattori di rischio significativi ed evitavano salti non supportati.

Significato e Rivendicazioni

Il paper sostiene che la supervisione CoT ancorata al KG consente ai piccoli LLM a pesi aperti di apprendere un ragionamento prognostico efficiente nei dati e clinicamente fondato. Il significato risiede nel colmare il divario tra predizione accurata e giustificazione interpretabile:

1. Affidabilità: Fondando il ragionamento in PrimeKG e filtrando per coerenza del label, il sistema mitiga le allucinazioni e garantisce la plausibilità medica.
2. Generalizzabilità: I pattern di ragionamento appresi si trasferiscono efficacemente a coorti fuori distribuzione (CRADLE), suggerendo che i modelli catturino i meccanismi sottostanti delle malattie piuttosto che artefatti specifici del dataset.
3. Utilità Clinica: Il framework produce spiegazioni che i clinici trovano più chiare e affidabili, facilitando potenzialmente l'adozione degli LLM nei sistemi di supporto alle decisioni cliniche dove la trasparenza è fondamentale.

Gli autori notano con moderazione che l'approccio dipende dalla copertura di PrimeKG e che i lavori futuri dovrebbero esplorare viste EHR più ricche (note, laboratori) e contesti multi-visita.