Tabular Foundation Models Can Do Survival Analysis

Sintesi Tecnica: I Modelli Fondazionali Tabulari Possono Effettuare l'Analisi della Sopravvivenza

Definizione del Problema
L'analisi della sopravvivenza, o analisi del tempo all'evento, mira a modellare quando si verifica un evento di interesse piuttosto che limitarsi a stabilire se esso avvenga. Una sfida determinante in questo dominio è la censura a destra (right-censoring), in cui il periodo di osservazione termina prima che l'evento si verifichi, lasciando sconosciuto il reale tempo dell'evento. Sebbene i Modelli Fondazionali Tabulari (TFM) abbiano ottenuto un successo straordinario nei compiti di classificazione e regressione — spesso superando metodi come XGBoost in regimi con dati limitati — adattarli all'analisi della sopravvivenza non è banale. Le architetture di regressione standard non possono gestire nativamente la censura, e i modelli esistenti specifici per la sopravvivenza richiedono spesso l'addestramento da zero su singoli dataset, limitandone l'utilità in contesti con pochi dati. Inoltre, mentre i TFM vengono valutati in termini di accuratezza o errore quadratico medio, la modellazione della sopravvivenza richiede probabilità ben calibrate per caratterizzare la dinamica del rischio nel tempo.

Metodologia
Gli autori propongono un framework che riformula l'analisi della sopravvivenza come una serie di problemi di classificazione binaria per sfruttare i TFM "off-the-shelf" tramite in-context learning, senza richiedere pre-addestramento specifico per la sopravvivenza, fine-tuning o modifiche architettoniche.

1. Discretizzazione Temporale: L'orizzonte temporale continuo viene partizionato in $K$ intervalli discreti.
2. Formulazione del Cedimento Cumulativo (Cumulative Failure): A differenza degli approoli tradizionali che modellano il tasso di rischio discreto (la probabilità che un evento si verifichi in un intervallo specifico dato che si è sopravvissuti fino a quel punto), questo metodo modella direttamente la probabilità di cedimento cumulativo. Per ogni punto temporale $t_k$, viene definito un'etichetta binaria $Y_{i,k} = 1(T_i \le t_k)$, che indica se l'evento si è verificato entro il tempo $t_k$.
3. Gestione della Censura: Le osservazioni censurate sono trattate come esempi con etichette mancanti per i punti temporali successivi al tempo di censura. L'obiettivo di addestramento è una standard perdita di cross-entropia binaria, limitata solo alle etichette osservate (dove $t_k < C_i$).
4. In-Context Learning: Il dataset riformulato (caratteristiche, punti temporali e etichette binarie) viene fornito a un TFM (ad esempio, MITRA, TabPFN) o a un classificatore standard (XGBoost). Il modello apprende a prevedere la probabilità di cedimento entro un tempo $t_k$ date le covariate.
5. Estensione Dinamica: Per contesti dinamici con covariate tempo-varianti, il modello prende la cronologia delle osservazioni fino al tempo $t_k$ e predice la probabilità di cedimento cumulativo per un orizzonte futuro $\Delta$, condizionata al fatto che il soggetto sia sopravvissuto fino al tempo $t_k$.

Garanzie Teoriche
Il documento fornisce una prova teorica (Teorema 4.2) dimostrando che, sotto l'ipotesi standard di censura condizionalmente indipendente, la minimizzazione della proposta perdita di cross-entropia binaria recupera le reali probabilità di sopravvivenza sottostanti man mano che la dimensione del set di addestramento tende all'infinito. Ciò stabilisce che la formulazione di classificazione è statisticamente coerente con l'obiettivo di sopravvivenza.

Contributi Chiave

1. Formulazione di Classificazione: Gli autori introducono una specifica formulazione di classificazione che modella direttamente le probabilità di cedimento cumulativo. Empiricamente, questa supera le formulazioni basate su tassi di rischio discreti (che soffrono di accumulo di errore moltiplicativo all'aumentare del numero di bin) e le formulazioni multiclasse.
2. Prova di Coerenza: Dimostrano che la minimizzazione della perdita di classificazione binaria produce stime coerenti delle probabilità di sopravvivenza sotto le normali assunzioni di censura.
3. Valutazione Empirica: L'approccio è valutato su 48 dataset reali (43 statici e 5 dinamici) dal repository SurvSet. Lo studio confronta i TFM off-the-shelf utilizzando questa formulazione contro i baseline classici (CoxPH, Random Survival Forests) e modelli di deep learning (DeepSurv, DeepHit, Joint Models).

Risultati

• Setting Statico: In 43 dataset, i TFM che utilizzano la formulazione di cedimento cumulativo (specificamente MITRA-CF e XGBoost-CF) hanno ottenuto le migliori prestazioni complessive su molteplici metriche (C-Index, AUC integrata, Brier Score integrato). La formulazione CF si è dimostrata più robusta rispetto al numero di bin di discretizzazione ($K$) rispetto agli approcci basati sul rischio (DH), che degradavano all'aumentare di $K$ a causa dell'accumulo di errore.
• Setting Dinamico: In 5 dataset dinamici, MITRA-CF ha superato tutti i baseline, inclusi CoxPH/RSF landmarked e Dynamic DeepHit. I risultati suggeriscono che i prior temporali incorporati nei TFM sono particolarmente efficaci per modellare le complessità dei dati longitudinali.
• Correlazione tra Metriche: Lo studio ha riscontrato una forte correlazione ($r=0,89$) tra la perdita di cross-entropia binaria e il Brier Score integrato, supportando l'affermazione teorica che l'ottimizzazione dell'obiettivo di classificazione migliora direttamente la calibrazione della probabilità di sopravvivenza.

Significatività e Rivendicazioni
Il documento afferma che le capacità di analisi della sopravvivenza possono essere elicitate da modelli fondazionali tabulari off-the-shelf attraverso una semplice riformulazione di classificazione binaria. Questo approccio offre un'alternativa computazionalmente efficiente all'addestramento di modelli specializzati per la sopravvivenza da zero, particolarmente nei regimi di piccoli dati dove i modelli fondazionali eccellono.

Gli autori distinguono il loro lavoro dagli sforzi concorrenti (Survival In-Context e SurvivalPFN) che pre-addestrano i transformer specificamente su compiti di sopravvivenza sintetici. Al contrario, questo lavoro dimostra che i TFM general-purpose, senza alcun pre-addestramento o modifica architettonica specifica per la sopravvivenza, possono raggiungere prestazioni allo stato dell'arte sui compiti di sopravvivenza quando forniti della corretta formulazione di classificazione. Il metodo è flessibile rispetto al numero di bin temporali ed è applicabile sia a contesti statici che dinamici, sebbene rimanga limitato a scenari che soddisfano l'assunzione di censura condizionalmente indipendente.