MedSAE: Dissecting MedCLIP Representations with Sparse Autoencoders

Riepilogo Tecnico: MEDSAE – Dissecting MedCLIP Representations with Sparse Autoencoders

Enunciato del Problema
L'integrazione dell'Intelligenza Artificiale (AI) nell'assistenza sanitaria, in particolare attraverso Modelli Linguistici Multimodali di grandi dimensioni (MLLM) come MedCLIP, ha migliorato le capacità diagnostiche ma introdotto sfide significative riguardo all'interpretabilità. Sebbene questi modelli raggiungano prestazioni elevate in compiti come l'analisi delle radiografie del torace, i loro meccanismi interni rimangono opachi "scatole nere". Operano in un regime di superposizione, dove il numero di caratteristiche apprese supera la dimensionalità dello spazio di rappresentazione, causando la moltiplicazione di più caratteristiche su neuroni condivisi. Ciò complica l'attribuzione di concetti clinici specifici ad attivazioni individuali del modello. Le tecniche esistenti di interpretabilità meccanicistica non sono state sufficientemente scalate o adattate per compiti complessi di visione-linguaggio medico, lasciando un divario tra l'AI medica ad alte prestazioni e la trasparenza richiesta per l'affidabilità clinica.

Metodologia
Gli autori propongono MedSAE, un framework che applica Autoencoder Sparsi (SAE) allo spazio latente di MedCLIP per disaccoppiare le caratteristiche sovrapposte in concetti comprensibili all'uomo, approssimativamente monosemantici. La pipeline consta di tre fasi principali:

1. Estrazione degli Embedding e Addestramento:

   • Gli autori utilizzano il codificatore visivo di MedCLIP (congelato) per generare embedding dal dataset CheXpert (radiografie del torace).
   • Per affrontare un eventuale disallineamento tra le modalità, gli embedding vengono normalizzati centrando (sottraendo le medie per modalità) e scalando per garantire norme coerenti.
   • Un SAE basato su ReLU viene addestrato su questi embedding normalizzati. Il SAE mappa vettori di input densi ($x$) in una rappresentazione latente sparsa ($z$) utilizzando un'architettura encoder-decoder con una penalità $L_1$ per imporre la sparsità:
     $$L := \|x - \hat{x}\|_2^2 + \lambda\|z\|_1$$
   • L'obiettivo è apprendere un insieme sovracomplesso di direzioni latenti che separino le caratteristiche moltiplicate.

2. Valutazione della Monosemanticità:

   • Per validare che i neuroni SAE corrispondano a singoli concetti, gli autori calcolano la correlazione di Pearson tra l'attivazione di ciascun neurone e le etichette cliniche codificate multi-hot (es. Edema, Polmonite).
   • Calcolano l'entropia della distribuzione delle correlazioni per ciascun neurone su tutte le etichette. Un'entropia più bassa indica che un neurone è fortemente concentrato su un singolo concetto (alta monosemanticità), mentre un'entropia più alta suggerisce un'attivazione distribuita su più concetti.

3. Assegnazione Automatica dei Nomi e Validazione:

   • Riconoscendo la difficoltà di costruire un vocabolario medico robusto per il matching basato su decoder, gli autori impiegano MedGemma, un modello fondazionale visione-linguaggio, per l'identificazione automatica dei concetti.
   • Per ciascun neurone SAE, le immagini con attivazione più alta vengono inviate a MedGemma tramite prompt strutturati per generare una spiegazione concettuale condivisa (es. "Edema polmonare severo").
   • Validazione: Per garantire che questi nomi generati non siano allucinazioni, viene eseguita un'attività di rilevamento. MedGemma viene presentato con la descrizione del concetto generata e un insieme bilanciato di immagini attivanti e non attivanti, quindi viene chiesto di classificare quali immagini corrispondono al concetto. Ciò fornisce una misura quantitativa dell'allineamento semantico.

Contributi Chiave

• Nuovo Framework di Valutazione: Il documento introduce un metodo di valutazione quantitativa che combina metriche basate sulla correlazione con la denominazione e la validazione automatica dei concetti utilizzando MedGemma. Questo supera la semplice analisi delle attivazioni fornendo etichette interpretabili dall'uomo.
• Validazione Empirica della Linearità: Gli autori dimostrano che gli embedding di MedCLIP soddisfano approssimativamente l'ipotesi di superposizione lineare. Attraverso un controllo di sanità mentale che coinvolge immagini composite, hanno mostrato che gli embedding composti possono essere ricostruiti linearmente dai loro componenti con un'alta similarità coseno (0.97), giustificando l'applicazione di strumenti di interpretabilità lineare come gli SAE.
• Disaccoppiamento delle Caratteristiche: Lo studio mostra che gli SAE possono estrarre con successo caratteristiche clinicamente significative e interpretabili dalle rappresentazioni interne di un modello visione-linguaggio medico, superando gli embedding grezzi di MedCLIP in termini di selettività di classe.

Risultati
Gli esperimenti condotti sul dataset CheXpert producono i seguenti risultati:

• Migliorata Monosemanticità: I neuroni SAE mostrano un'entropia inferiore (2.25) rispetto agli embedding grezzi di MedCLIP (2.38), indicando un allineamento più forte con singoli concetti.
• Maggiore Accuratezza Interpretativa: Il framework di denominazione automatica ha identificato 21 neuroni SAE distinti con un'accuratezza di rilevamento superiore al 70% (es. "Versamento pleurico destro con possibili loculazioni", "Cardiomegalia severa con congestione vascolare polmonare significativa"). Al contrario, gli embedding grezzi di MedCLIP hanno prodotto solo due interpretazioni ad alta accuratezza di questo tipo.
• Robustezza dei Risultati: Una condizione di controllo in cui le coppie neurone-concetto sono state mescolate casualmente ha portato a un calo significativo delle prestazioni di rilevamento, confermando che le interpretazioni significative derivano dalla struttura SAE e non da artefatti statistici.
• Bilanciamento Sparsità-Ricostruzione: Il modello selezionato ha raggiunto un equilibrio favorevole con un tasso di attivazione $L_0$ dello 0.20% e una frazione di varianza spiegata (FVE) di 0.98, nonostante un tasso del 30% di neuroni "morti".

Significato e Affermazioni
Il documento posiziona MedSAE come un passo scalabile verso rappresentazioni clinicamente affidabili nell'AI medica. Colmando il divario tra modelli ad alte prestazioni e trasparenza, il lavoro offre una via per invertire l'ingegneria dei meccanismi computazionali in concetti medici comprensibili all'uomo. Gli autori affermano che questa è la prima analisi empirica dell'interpretabilità meccanicistica negli MLLM medici utilizzando SAE.

Tuttavia, gli autori mantengono un tono modesto riguardo alle limitazioni. Riconoscono che il loro approccio si basa sull'ipotesi di rappresentazione lineare, che potrebbe non valere universalmente, e che la loro analisi attuale è limitata a un singolo livello (MedCLIP-ResNet). Osservano che le unità inattive potrebbero riflettere tensioni tra capacità e sparsità piuttosto che un fallimento dell'approccio e che la validazione con intervento umano da parte di esperti medici rimane necessaria per valutare la rilevanza clinica e la sicurezza. Il lavoro è presentato come un passo fondazionale per motivare ulteriori ricerche su sistemi di AI medica trasparenti, robusti e affidabili.