Efficient and Effective Internal Memory Retrieval for LLM-Based Healthcare Prediction

1. Il Problema

L'uso dei Modelli Linguistici di Grande Dimensione (LLM) in ambito sanitario promette molto, ma la loro applicazione in contesti clinici ad alto rischio è ostacolata da due fattori principali:

• Allucinazioni e mancanza di contesto: Gli LLM tendono a generare informazioni non veritiere e faticano ad accedere a contesti medici granulari e specifici per il paziente.
• Limiti delle soluzioni attuali (RAG): L'approccio standard, la Generazione Aumentata dal Recupero (RAG), richiede la ricerca in enormi basi di conoscenza esterne. Questo introduce:
  • Alta latenza: Le ricerche esterne sono computazionalmente costose, rendendo il sistema impraticabile per cure che richiedono decisioni rapide.
  • Costi di contesto: L'iniezione di conoscenze esterne tramite prompt allunga il contesto di input, aumentando i costi di inferenza e limitando la scalabilità.
  • Difficoltà di recupero: Costruire recuperatori supervisionati di alta qualità richiede coppie query-contesto etichettate estese o ricerche su grafi complesse.

2. Metodologia: Il Framework K2K

Per superare questi limiti, gli autori introducono K2K (Keys-to-Knowledge), un framework che bypassa il recupero esterno recuperando direttamente la conoscenza dalla memoria interna (spazio dei parametri) dell'LLM. Il framework si basa sull'interpretazione delle reti Feed-Forward (FFN) nei transformer come sistemi di memoria associativa chiave-valore.

Il processo avviene in tre fasi principali (illustrate nella Figura 1 e 2 del paper):

A. Costruzione della Memoria di Recupero (Internal Memory Construction)

Invece di cercare documenti esterni, K2K trasforma l'LLM pre-addestrato in una memoria di recupero interna:

• Memoria a Livello di Documento: Utilizza un modello adattato al dominio (tramite LoRA) per iniettare conoscenze mediche testuali. Le matrici dei pesi $W_1$ (chiavi) degli strati FFN diventano la rappresentazione interna della conoscenza documentale.
• Memoria a Livello di Grafo: Le triple di un grafo della conoscenza medica vengono linearizzate in testo e iniettate nel modello tramite LoRA. Le matrici adattate $A_1B_1$ rappresentano la conoscenza strutturata interna.
• Vantaggio: Questo elimina la necessità di cercare in basi dati esterne durante l'inferenza.

B. Costruzione della Query Guida dall'Attivazione (Activation-Guided Probe Construction)

Un semplice pooling medio delle query non è sufficiente per recuperare chiavi rilevanti, poiché le rappresentazioni tendono a essere poco discriminative. K2K introduce un meccanismo avanzato:

• Identificazione dei Token Critici: Calcola un "bias di attivazione" per identificare i token più informativi e le caratteristiche "outlier" all'interno di una finestra temporale.
• Distanza di Mahalanobis: Invece della distanza euclidea (che tratta tutte le dimensioni come uguali), K2K utilizza una approssimazione diagonale della distanza di Mahalanobis. Questo pesa le deviazioni token per token considerando la varianza per dimensione, enfatizzando le direzioni a bassa varianza che sono spesso semanticamente più discriminative.
• Risultato: Si ottiene una vettore di sonda (probe) potenziato ($Q^+$) che cattura meglio le sfumature semantiche del contesto clinico.

C. Riperclassificazione con Attenzione Incrociata (Cross-Attentive Reranking)

Una volta recuperate le $k$ chiavi più simili dalla memoria interna (documentale e grafica):

• Integrazione Dinamica: Un meccanismo di attenzione incrociata (Cross-Attention) riperclassifica e fonde le conoscenze recuperate in base alla query specifica.
• Fusione: Le conoscenze di documento e grafo vengono normalizzate, concatenate e fuse con le rappresentazioni originali dell'input per la previsione finale tramite un MLP.

3. Contributi Chiave

1. Recupero Interno vs. Esterno: Sostituzione del recupero esterno (RAG) con l'accesso diretto alle chiavi interne dell'LLM, eliminando la latenza di ricerca e i costi di contesto esteso.
2. Iniezione di Conoscenza tramite LoRA: Utilizzo efficiente di Low-Rank Adaptation per infondere conoscenze mediche specifiche (testuali e strutturate) direttamente nello spazio dei parametri del modello senza ri-addestramento completo.
3. Query Guidata dall'Attivazione: Introduzione di un metodo di costruzione della query basato sulla distanza di Mahalanobis per migliorare la discriminazione semantica e la precisione del recupero.
4. Riperclassificazione Adattiva: Un modulo di Cross-Attention che integra dinamicamente fonti multiple di conoscenza interna in base al compito specifico.

4. Risultati Sperimentali

Il framework è stato valutato su quattro dataset derivati da MIMIC-III e MIMIC-IV, focalizzandosi su due task critici: Previsione di Mortalità e Previsione di Riammissione.

• Performance: K2K ha raggiunto lo stato dell'arte (SOTA) su tutti i dataset, superando modelli sequenziali classici (GRU, RETAIN), modelli basati su recupero esterno (KARE, RAG standard) e modelli generativi basati su prompt.
  • Su Mortality-MIMIC-III, K2K ha ottenuto un punteggio medio (F1, Jaccard, AUPRC, AUROC) di 26.26, superando KARE (23.57) e RAG standard (20.45).
  • Su Readmission-MIMIC-III, K2K ha ottenuto un punteggio medio di 60.37, contro il 53.46 di KARE.
• Efficienza: K2K è significativamente più veloce in termini di tempo di recupero rispetto ai metodi basati su RAG esterno, rendendolo adatto a scenari clinici in tempo reale.
• Analisi Ablative:
  • La rimozione di una delle fonti di conoscenza (documentale o grafo) riduce le prestazioni, dimostrando che l'integrazione di entrambe è cruciale.
  • L'uso della distanza di Mahalanobis per la costruzione della query supera nettamente l'uso della distanza euclidea o del semplice pooling medio.
  • L'estrazione di conoscenza da diversi strati del transformer mostra che sia gli strati superficiali che quelli profondi contengono informazioni preziose, sfatando il mito che solo gli strati superiori siano utili per la semantica.

5. Significato e Implicazioni

Questo lavoro rappresenta un cambio di paradigma nell'uso degli LLM per la sanità:

• Superamento del collo di bottiglia computazionale: Dimostra che la conoscenza non deve necessariamente essere cercata all'esterno; può essere "estratta" efficientemente dai parametri stessi del modello se opportunamente strutturata e interrogata.
• Affidabilità Clinica: Riducendo le allucinazioni e fornendo un recupero contestuale preciso e a bassa latenza, K2K aumenta la fiducia nell'uso degli LLM per decisioni cliniche critiche.
• Scalabilità: L'approccio interno rende i sistemi di supporto decisionale clinico più scalabili ed economici, eliminando la dipendenza da infrastrutture di ricerca esterne massive.

In sintesi, K2K trasforma l'LLM da un semplice generatore di testo in un sistema di ragionamento clinico autonomo, capace di accedere alla propria "memoria interna" in modo rapido e preciso, offrendo una soluzione superiore ai tradizionali approcci RAG per le previsioni sanitarie.