Structure-Aware Set Transformers: Temporal and Variable-Type Attention Biases for Asynchronous Clinical Time Series

Il paper presenta STAR Set Transformer, un modello che migliora l'analisi delle serie temporali cliniche asincrone introducendo bias di attenzione temporale e di tipo variabile per preservare la struttura dei dati senza ricorrere all'imputazione, ottenendo prestazioni superiori su compiti predittivi in terapia intensiva.

L'essenziale

🏥 Il Problema: La Cartella Clinica "Disordinata"

Immagina di dover leggere la cartella clinica di un paziente in terapia intensiva. Non è come leggere un libro, dove le parole sono in ordine e su righe fisse. È più come avere un mucchio di post-it sparsi sul pavimento:

• Alcuni post-it sono scritti alle 3:00, altri alle 3:15, altri alle 4:30.
• Alcuni parlano della febbre, altri della pressione, altri del battito cardiaco.
• Spesso mancano pezzi: non c'è un post-it per la febbre alle 3:15, ma ce n'è uno per la pressione.

I computer, quando imparano a leggere questi dati (chiamati serie temporali), hanno due modi principali per organizzarli:

1. La griglia rigida (Il Calendario): Costringi tutto in un calendario orario (ogni ora ha una riga). Se manca un dato, devi inventartelo (imputazione) o mettere una "spunta" che dice "qui manca qualcosa". Il rischio? Il computer potrebbe imparare a fare previsioni basandosi solo su quanto manca, e non su cosa manca davvero.
2. La lista disordinata (Il Set): Prendi ogni singolo post-it e lo metti in una lista. È fedele alla realtà, ma il computer perde la "geografia": non vede più che la febbre e la pressione sono state misurate quasi insieme, né che la febbre di oggi è collegata a quella di ieri.

💡 La Soluzione: Gli "Occhiali Magici" (STAR)

Gli autori di questo paper hanno creato un nuovo modello chiamato STAR (Structure-AwaRe Set Transformer). Invece di forzare i dati in una griglia o lasciarli completamente disordinati, hanno dato al computer degli "occhiali magici" (chiamati bias di attenzione) che gli permettono di vedere le connessioni nascoste, anche se i dati sono sparsi.

Ecco come funzionano questi occhiali, usando due metafore semplici:

1. L'Occhiale del "Vicinato Temporale" (Temporal Bias)

Immagina che ogni post-it sia una persona in una stanza.

• Senza occhiali: Il computer guarda tutti i post-it con la stessa probabilità, come se tutti fossero ugualmente importanti, indipendentemente da quando sono stati scritti.
• Con l'occhiale temporale: Il computer impara che le persone che sono state nella stanza poco fa (post-it vicini nel tempo) hanno più probabilità di parlarsi tra loro rispetto a quelle che erano lì due giorni fa.
  • Esempio: Se il battito cardiaco sale alle 10:00, il computer capisce che è più probabile che sia collegato alla pressione misurata alle 10:05 piuttosto che alla febbre misurata alle 8:00. È come dire: "Chi è vicino nel tempo, è vicino anche nel ragionamento".

2. L'Occhiale del "Gruppo di Appartenenza" (Variable-Type Bias)

Ora immagina che ogni post-it abbia un adesivo colorato che indica di cosa parla (Rosso = Cuore, Blu = Polmoni, Verde = Reni).

• Senza occhiali: Il computer potrebbe confondersi e pensare che un dato sui reni sia collegato a un dato sul cuore solo perché sono vicini nel tempo.
• Con l'occhiale di tipo: Il computer impara che i post-it rossi (cuore) dovrebbero "parlare" più spesso tra loro, e i blu (polmoni) tra loro.
  • Esempio: Se il computer vede un dato sulla glicemia (Rosso) e uno sulla pressione (Rosso), capisce che c'è una forte affinità, anche se sono stati misurati in momenti leggermente diversi. È come se il computer dicesse: "Aspetta, questi due dati appartengono alla stessa famiglia, quindi devo ascoltarli insieme".

🚀 Come hanno testato la magia?

Gli autori hanno fatto un esperimento curioso: hanno provato a mettere questi "occhiali" in diverse parti del cervello del computer (gli strati del Transformer), come se provassero a insegnare la regola del "vicinato" o del "gruppo" prima o dopo che il computer aveva già iniziato a pensare.

Hanno scoperto che:

• Se dai al computer entrambi gli occhiali (sia il tempo che il tipo di dato) in modo intelligente, funziona meglio di chiunque altro.
• Il modello STAR ha battuto tutti i record su tre compiti critici:
  1. Prevedere un arresto cardiaco (CPR).
  2. Prevedere il rischio di morte (Mortalità).
  3. Prevedere la necessità di farmaci per la pressione (Vasopressori).

🌟 Perché è importante?

Fino a ora, per far capire ai computer i dati medici, dovevamo scegliere tra:

• Rigidità: Costruire tabelle perfette ma perdere informazioni reali.
• Caos: Lasciare i dati liberi ma perdere il contesto.

STAR è come un traduttore perfetto che prende il caos dei post-it sparsi e dice al computer: "Ehi, guarda che anche se non sono in una tabella, questi due dati sono vicini nel tempo e parlano della stessa cosa. Mettili in relazione!".

In pratica, hanno reso l'intelligenza artificiale più "umana" nel leggere le cartelle cliniche: non si basa su regole rigide, ma capisce il contesto, il tempo e le relazioni tra le cose, proprio come farebbe un medico esperto. E i risultati? Molto più precisi nel salvare vite umane.

Sommario tecnico

1. Il Problema: La Rappresentazione dei Dati EHR

Le cartelle cliniche elettroniche (EHR) sono serie temporali multivariate irregolari e asincrone. A differenza di immagini o testo, non esiste una discretizzazione canonica del tempo, rendendo la scelta del layout di input cruciale per i modelli neurali.

Il paper identifica tre approcci principali e i loro limiti:

• Griglie Regolari (Regular Grids): Il tempo è binnato in intervalli fissi (es. orari). Richiede imputazione dei dati mancanti o l'uso di maschere di missingness, rischiando di introdurre errori o far sì che il modello impari scorciatoie basate sulle politiche di campionamento (es. "se un dato manca, il paziente è stabile").
• Griglie basate sugli Eventi (Event-Time Grids): Indicano il tempo unendo tutti i timestamp osservati. Sebbene evitino l'imputazione fittizia, l'asincronia genera ancora molta missingness, richiedendo maschere o feature di "gap temporale".
• Tokenizzazione a Insiemi (Point-Set Tokenization): Tratta ogni evento osservato come un token (valore, tempo, identità della variabile) senza discretizzare il tempo. Sebbene eviti l'imputazione, perde due assi strutturali fondamentali:
  1. Le traiettorie temporali entro una singola variabile (colonne).
  2. Le relazioni cross-variabili contemporanee (righe).
     In questo approccio, l'attenzione deve recuperare queste strutture dai dati da sola, senza indizi induttivi (priors).

2. Metodologia: STAR Set Transformer

Gli autori propongono STAR (Structure-AwaRe) Set Transformer, un modello che combina la flessibilità della tokenizzazione a insiemi con l'efficienza dei bias di attenzione "soft" per reintrodurre le strutture perse.

Componenti Chiave:

1. Input: Un episodio EHR è rappresentato come un insieme di triplette $(t, v, s)$: timestamp, valore osservato e indice della variabile. Vengono aggiunti due token speciali: un token [CLS] e un token demografico (età/sesso).
2. Embedder: Utilizza un embedder di insiemi (basato su ITE) per generare una sequenza di token iniziali.
3. Bias di Attenzione Additivi: La modifica core avviene nei logit dell'attenzione del Transformer. Invece di affidarsi solo al prodotto scalare contenuto, vengono aggiunti due bias parametrici ed efficienti:
   • Bias Temporale (Temporal Bias): Introduce una penalità basata sulla distanza temporale tra due token $i$ e $j$.
     $$b_{time} = - \frac{|t_i - t_j|}{\tau}$$
     Dove $\tau$ è una scala temporale apprendibile (parametrizzata come $\exp(\omega) + \epsilon$). Questo favorisce le interazioni tra token temporalmente vicini.
   • Bias Tipo-Variabile (Variable-Type Bias): Utilizza una matrice di affinità appresa $B$ tra i tipi di variabili.
     $$b_{var} = B_{s_i, s_j}$$
     Questo favorisce le interazioni tra token della stessa variabile (o variabili compatibili), ripristinando la struttura "colonna" delle serie temporali.

Strategia di Fusione a Livello di Strato (Layer-wise Fusion):

Il paper esplora sistematicamente dove e quando iniettare questi bias all'interno di un Transformer a 4 strati. Vengono testati 10 diversi schemi di fusione (es. bias temporale solo negli strati bassi, bias di tipo solo in quelli alti, o combinazioni miste). La configurazione proposta come migliore è vt-vt, dove entrambi i bias sono attivi in tutti gli strati.

3. Contributi Principali

1. Biasing dell'Attenzione per EHR Irregolari: Introduzione del STAR Set Transformer che recupera la struttura induttiva delle griglie (tempo e variabile) senza dover discretizzare il tempo, mantenendo la flessibilità degli insiemi.
2. Due Bias Complementari ed Efficienti:
   • Un bias temporale basato su una penalità di distanza appresa.
   • Un bias di tipo variabile basato su una matrice di compatibilità appresa.
3. Ablazione Sistematica: Uno studio approfondito sulle strategie di iniezione dei bias attraverso la profondità del Transformer, identificando che la combinazione di entrambi i bias in tutti gli strati (o in schemi specifici) porta a miglioramenti coerenti.

4. Risultati Sperimentali

Il modello è stato valutato su tre compiti di previsione clinica utilizzando il dataset MIMIC-IV:

• Rianimazione Cardiopolmonare (CPR)
• Mortalità
• Uso di Vasopressori

Performance (AUC / APR):
STAR Set Transformer ha superato tutti i baseline, inclusi modelli basati su griglie regolari (SMART, DueTT), griglie a eventi (PrimeNet) e approcci a insiemi precedenti (STraTS).

Task	Modello	AUC	APR
CPR	STAR (Ours)	0.7158	0.0026
	DueTT (Grid)	0.6478	0.0010
	STraTS (Set)	0.5397	0.0018
Mortalità	STAR (Ours)	0.9164	0.2033
	STraTS (Set)	0.8778	0.1457
Vasopressori	STAR (Ours)	0.8373	0.1258
	DueTT (Grid)	0.8255	0.0929

Analisi degli Ablation:

• Il bias temporale è stato il principale motore dei guadagni in AUC (specialmente per CPR, passando da 0.600 a 0.753 rispetto al baseline senza bias).
• Il bias di tipo variabile ha fornito miglioramenti consistenti ma più piccoli.
• La combinazione di entrambi (vt-vt) ha ottenuto il miglior APR medio (0.111) e le migliori performance su Mortalità e Vasopressori, dimostrando che la struttura induttiva aiuta sia la discriminazione che la calibrazione delle probabilità.

5. Significato e Conclusioni

Il lavoro dimostra che è possibile ottenere i vantaggi strutturali delle griglie temporali (traiettorie chiare e relazioni cross-variabili) senza i difetti dell'imputazione o della discretizzazione forzata.

• Interpretabilità: I parametri appresi ($\tau$ e la matrice $B$) forniscono riassunti interpretabili delle scale temporali contestuali e delle interazioni tra variabili.
• Praticità: L'approccio è un "plug-in" efficiente per modelli di serie temporali basati su trasformatori, che non richiede cambiamenti architetturali massicci ma solo l'aggiunta di bias additivi.
• Impatto: Offre una soluzione robusta al problema dell'asincronia nei dati clinici, riducendo la dipendenza da scorciatoie apprese dal modello basate sui pattern di missingness.

In sintesi, STAR Set Transformer rappresenta un passo avanti significativo nell'elaborazione di serie temporali cliniche irregolari, bilanciando la flessibilità dei dati grezzi con la necessità di induttività strutturale per modelli di fondazione (foundation models) efficaci.