Distributionally Robust Transfer Learning with Structurally Missing Covariates, with Application to Cross-National Cardiac Arrest Prediction

Sintesi Tecnica: Apprendimento per Trasferimento Robusto Distribuzionalmente con Covariate Strutturalmente Mancanti

Enunciato del Problema
Il documento affronta una sfida critica nel dispiegamento di modelli predittivi clinici attraverso diversi sistemi sanitari: lo scenario in cui un modello addestrato in un contesto "sorgente" ad alta risorsa si basa su un ricco insieme di covariate $(X, A)$, ma il contesto di dispiegamento "target" manca completamente di specifiche covariate altamente predittive $A$ a causa di differenze strutturali nell'infrastruttura di raccolta dati. A differenza dei dati mancanti standard, dove le osservazioni individuali sono incomplete, questa "mancanza strutturale" rappresenta un'assenza a livello di popolazione di variabili. Inoltre, il dominio target spesso manca di dati di outcome etichettati, precludendo l'adattamento supervisionato o il retraining.

Gli approcci esistenti falliscono in questo contesto. I metodi di imputazione richiedono assunzioni non verificabili sulla distribuzione di $A$ nel dominio target. Le correzioni per lo spostamento delle covariate (ad esempio, la ponderazione per importanza) affrontano solo le differenze nelle variabili osservate $X$ e non possono recuperare le informazioni perse dall'assenza di $A$. L'ottimizzazione robusta distribuzionalmente (DRO) standard costruisce tipicamente insiemi di incertezza sulla distribuzione congiunta di tutte le variabili, il che è eccessivamente conservativo e non riesce a distinguere tra covariate osservate stabili e quelle strutturalmente mancanti.

Metodologia: Il Framework DRUM
Gli autori propongono DRUM (Distributionally Robust Unsupervised transfer learning with structurally Missing covariates), un framework progettato per trasferire modelli predittivi a popolazioni target non etichettate dove $A$ è assente. La filosofia di fondo è ottimizzare le prestazioni predittive nel caso peggiore su un insieme di incertezza della distribuzione condizionale target sconosciuta $A|X$, piuttosto che tentare di imputare $A$ o assumere che $A|X$ rimanga stabile.

La metodologia procede in tre fasi:

1. Formulazione del Problema:
   L'obiettivo è apprendere un predittore $m(X)$ utilizzando solo le covariate condivise $X$. Gli autori formulano questo come un problema minimax:
   $$\min_{m(\cdot)} \max_{P_{A|X} \in \mathcal{C}(\delta)} \mathbb{E}_{X \sim Q_X, A \sim P_{A|X}, Y \sim P^S_{Y|X,A}} \left[ Y^2 - (Y - m(X))^2 \right]$$
   Qui, $\mathcal{C}(\delta)$ è un insieme di incertezza di distribuzioni condizionali per $A|X$ nel dominio target, controllato da un parametro di robustezza $\delta$. L'obiettivo massimizza la riduzione nel caso peggiore dell'errore quadratico rispetto a un predittore nullo.

2. Parametrizzazione Basata su Generatore e DRO:
   L'ottimizzazione diretta sullo spazio infinito-dimensionale delle distribuzioni condizionali è intrattabile. DRUM parametrizza la distribuzione condizionale $A|X$ utilizzando un generatore di rete neurale $g_\phi(X, \epsilon)$, dove $\epsilon$ è rumore latente.

   • Insieme di Incertezza: L'insieme di incertezza è definito utilizzando la Distanza Energetica (tramite il framework Engression). L'insieme $\mathcal{C}_g(\delta)$ contiene generatori il cui punteggio energetico devia dal generatore condizionale sorgente $g^S$ di al massimo $\delta$. Ciò permette al modello di esplorare distribuzioni condizionali che si discostano dalla sorgente senza richiedere la stima della densità.
   • Ottimizzazione: Il problema è risolto tramite un approccio primale-duale. Un generatore "nel caso peggiore" viene addestrato per minimizzare l'errore quadratico di previsione atteso sul target $X$ (marginalizzando su $A$ generato), soggetto al vincolo della distanza energetica.

3. Correzione del Bias (Ortogonalità di Neyman):
   Gli autori identificano che lo stimatore "plug-in" (mediante la media del modello di outcome sorgente $\hat{f}(X, A)$ su $A$ generato) eredita bias del primo ordine dalla stima di $\hat{f}$. Per mitigare ciò, derivano uno stimatore corretto per il bias basato sulla stima ortogonale di Neyman e sul cross-fitting.

   • Costruiscono un pseudo-outcome $F_i$ che combina i residui sorgente (ripesati da un rapporto di densità $\hat{\omega}$ per corrispondere alla distribuzione indotta dal generatore) e le previsioni plug-in target.
   • La regressione di questo pseudo-outcome su $X$ produce uno stimatore robusto agli errori del primo ordine nelle funzioni di disturbo ($\hat{f}$ e il rapporto di densità $\hat{\omega}$), raggiungendo tassi di convergenza parametrica sotto condizioni di prodotto dei tassi.

Contributi Chiave

1. Un Nuovo Framework DRO per la Mancanza Strutturale: DRUM isola l'incertezza specificamente alla distribuzione condizionale $A|X$ piuttosto che alla distribuzione congiunta, utilizzando un modello generativo per parametrizzare l'insieme di incertezza. Ciò evita l'intrattabilità dell'ottimizzazione infinito-dimensionale.
2. Stima Corretta per il Bias: Il documento deriva uno stimatore debiased che riduce la sensibilità agli errori di stima delle funzioni di disturbo, migliorando significativamente le prestazioni su campioni finiti nelle simulazioni e nelle applicazioni reali.
3. Applicazione alla Predizione OHCA Transnazionale: Il framework è applicato alla previsione degli esiti neurologici dopo arresto cardiaco extra-ospedaliero (OHCA), trasferendo modelli dal registro US-ROC (ricchi dati pre-ospedalieri) ai registri asiatici PAROS (dove variabili pre-ospedaliere come dose di epinefrina e pH del sangue sono strutturalmente mancanti).

Risultati

• Simulazioni: In contesti con relazioni condizionali lineari e non lineari e dimensioni variabili delle covariate mancanti, DRUM (in particolare le varianti debiased) ha superato sostanzialmente le linee di base standard (ERM, DRO, ponderazione per importanza) e i metodi di imputazione con pseudo-label. Ha raggiunto un Errore Quadratico Medio (MSE) normalizzato medio e nel caso peggiore inferiore sotto diversi gradi di spostamento distribuzionale.
• Applicazione Reale (OHCA):
  • Calibrazione: Quando applicato a Singapore, Giappone e Corea, DRUM ha prodotto previsioni significativamente meglio calibrate rispetto alle linee di base. Le linee di base standard sovrastimavano sistematicamente gli esiti favorevoli a causa dello spostamento di prevalenza e delle covariate mancanti, mentre DRUM seguiva più da vicino i tassi di evento osservati.
  • Prestazioni: Le varianti DRUM-Debiased hanno ottenuto i punteggi Brier più bassi e l'Errore di Calibrazione Atteso (ECE) più basso in tutti i siti.
  • Utilità Clinica: A soglie di probabilità fisse rilevanti per il triage clinico, DRUM-Debiased ha dimostrato maggiore precisione e specificità, indicando che le sue probabilità calibrate portano a classificazioni positive più affidabili rispetto alle linee di base che soffrivano di alti tassi di falsi positivi a causa di una scarsa calibrazione.

Significato e Affermazioni
Il documento afferma che DRUM fornisce un'alternativa principiale alla dicotomia attuale nel dispiegamento di AI clinica: o limitare i modelli a siti con infrastruttura dati identica o scartare informazioni predittive ignorando le covariate mancanti. Tenendo conto esplicitamente dell'incertezza nella distribuzione condizionale non osservata $A|X$ senza fare affidamento su assunzioni non verificabili, DRUM consente il dispiegamento di modelli ad alta risorsa in contesti a bassa risorsa.

Gli autori sottolineano che il vantaggio principale del metodo risiede nella calibrazione e nella stratificazione del rischio (punteggio Brier, ECE, precisione a soglia fissa) piuttosto che nella discriminazione globale (AUROC), il che è coerente con l'obiettivo del caso peggiore che penalizza l'errore quadratico. Il framework è presentato come una soluzione generale per qualsiasi problema di apprendimento per trasferimento in cui le covariate predittive sono disponibili durante l'addestramento ma assenti o limitate al dispiegamento, con potenziali estensioni alla equità algoritmica e alla previsione della stabilità farmaceutica.

Il documento rimane modesto riguardo alle limitazioni, notando che la parametrizzazione del generatore attuale assume covariate continue, il framework gestisce attualmente una singola popolazione sorgente e la derivazione si concentra sulle perdite di errore quadratico e cross-entropy. Le direzioni future includono la gestione di variabili discrete/miste, l'apprendimento federato con più sorgenti e outcome strutturati come dati tempo-evento.