ImmuVis: Hyperconvolutional Foundation Model for Imaging Mass Cytometry

Riepilogo Tecnico: ImmuVis

Enunciato del Problema
La Citometria di Massa per Imaging (IMC) consente il profilamento spaziale ad alto rendimento e a livello subcellulare dei tessuti mediante la misurazione di anticorpi marcati con metalli. Tuttavia, una sfida fondamentale nella modellazione dei dati IMC è l'assenza di uno spazio di canali fissi. A differenza delle immagini naturali con un set statico di canali RGB, i pannelli IMC variano significativamente tra gli studi a causa delle interferenze fisiche dei canali e dei disegni sperimentali specifici per coorte. Questa eterogeneità viola l'assunzione fondamentale dei backbones di visione standard, che richiedono una dimensione di input fissa. Gli approcci esistenti, come le U-Net, sono limitati a pannelli di marcatori fissi, mentre i recenti modelli basati su Transformer (ad esempio VirTues, EVA) tentano di gestire input variabili concatenando token spaziali specifici per marcatore con codifiche di identità. Questi metodi basati su token spesso ignorano le dipendenze locali cruciali tra i marcatori e comportano costi computazionali elevati all'aumentare del numero di marcatori e delle risoluzioni spaziali. Inoltre, la colorazione virtuale (previsione di marcatori non misurati) è un problema mal posto quando le evidenze sono deboli; i modelli attuali spesso mancano di stime di incertezza calibrate, rendendo difficile distinguere tra ricostruzioni biologicamente plausibili e previsioni errate.

Metodologia
Gli autori introducono ImmuVis, una famiglia di modelli fondazione progettata per unificare l'apprendimento con input variabili, l'efficienza computazionale e la quantificazione dell'incertezza. L'innovazione centrale è un'architettura di iperconvoluzione adattiva ai marcatori.

1. Architettura Iperconvoluzionale:

   • Codifica Indipendente dai Marcatori: I canali di input dei marcatori vengono inizialmente elaborati in modo indipendente da un tronco convoluzionale condiviso ($Enc_{ma}$) per catturare pattern spaziali locali.
   • Iperconvoluzione Condizionale: Invece di pesi fissi, il modello impiega una iperrete che genera kernel convoluzionali condizionati da codifiche di marcatori apprese. Per un insieme di marcatori di input $I_e$, un iperkernel $H^e_{I_e}$ viene costruito concatenando generatori di kernel appresi $\phi_e(i)$ per ogni marcatore $i \in I_e$. Ciò permette al modello di adattare dinamicamente i propri operatori di estrazione delle caratteristiche al sottoinsieme specifico di marcatori osservati senza modifiche architetturali.
   • Backbone Pan-Marcatore: L'output dell'iperconvoluzione è un embedding unificato e di dimensione fissa che può essere elaborato da backbone standard di visione artificiale (ad esempio ConvNeXt v2, Swin Transformer, ViT).
   • Decodificatore Simmetrico: Il decodificatore rispecchia il codificatore, utilizzando un'iperconvoluzione condizionata ai marcatori per mappare la rappresentazione latente su specifici marcatori target, prevedendo sia l'espressione media che la log-varianza (incertezza).

2. Strategia di Addestramento:

   • Dataset: I modelli sono pre-addestrati su IMC17M, il più grande dataset IMC a oggi, comprendente 24.405 immagini, 265 marcatori e oltre 17 milioni di patch attraverso 28 coorti.
   • Obiettivo: I modelli sono addestrati utilizzando un'attività di ricostruzione di canali mascherati con una funzione di perdita di regressione eteroscedastica Gaussiana. Questo obiettivo modella i pixel target come distribuzioni Gaussiane indipendenti, permettendo alla rete di prevedere sia la media ($\mu$) che la varianza ($\sigma^2$). La previsione della varianza fornisce una stima di incertezza calibrata che correla con l'errore di ricostruzione.
   • Preprocessing: Gli autori impiegano una specifica pipeline di preprocessing che coinvolge la trasformazione arcsinh (per gestire la dinamica del conteggio degli ioni), il filtraggio passa-basso di Butterworth (per ridurre il rumore specifico dell'IMC) e la scalatura al 99-esimo percentile minimo.

Contributi Chiave

• Design Iperconvoluzionale: Un'architettura innovativa che genera kernel convoluzionali da codifiche di marcatori apprese, permettendo a un singolo modello di elaborare sottoinsiemi arbitrari di marcatori senza re-addestramento o modifica architetturale.
• Famiglia Efficiente di Modelli Fondazione: Implementazione di tre varianti (ImmuVisConv, ImmuVisSwin, ImmuVisViT) che sfruttano backbone consolidati. Questi modelli raggiungono prestazioni all'avanguardia con costi computazionali significativamente inferiori rispetto alle alternative Transformer basate sulla concatenazione di token.
• Apprendimento Consapevole dell'Incertezza: Integrazione di un obiettivo di verosimiglianza eteroscedastica che produce stime di incertezza pixel per pixel, affrontando le questioni di affidabilità intrinseche nella colorazione virtuale per fenotipi rari o pattern atipici.
• Pre-addestramento su Grande Scala: La curatela e l'utilizzo di IMC17M, facilitando l'apprendimento di rappresentazioni robuste e trasferibili attraverso diverse istologie e pannelli di marcatori.

Risultati

• Colorazione Virtuale: Le varianti di ImmuVis superano le basi all'avanguardia (VirTues, EVA) nelle attività di colorazione virtuale. Nello specifico, ImmuVisSwin e ImmuVisConv hanno ottenuto il più basso Errore Quadratico Medio (MSE) e la più alta correlazione di Pearson, dimostrando una conservazione superiore della struttura spaziale rispetto ai metodi basati su token.
• Calibrazione dell'Incertezza: La testata eteroscedastica produce stime di incertezza che correlano fortemente con l'errore di ricostruzione (correlazione di Pearson > 0,92). I modelli forniscono segnali di confidenza significativi per pixel, con ImmuVisSwin che ottiene il più basso Area Under Sparsification Error (AUSE).
• Apprendimento delle Rappresentazioni: Nelle attività a valle, gli embedding di ImmuVis hanno prodotto prestazioni all'avanguardia nel tipaggio cellulare (ImmuVisConv ha ottenuto un Macro-F1 di 0,709) e nella previsione clinica (ImmuVisSwin ha ottenuto un Macro-F1 di 0,641), superando le varianti di VirTues ed EVA.
• Efficienza Computazionale: I modelli ImmuVis operano a un costo computazionale sostanzialmente inferiore. ImmuVisConv richiede circa 89,9 GFLOPs per passaggio in avanti, che è un ordine di grandezza inferiore rispetto a VirTues (1049,3 GFLOPs) e quasi 40 volte inferiore rispetto a EVA (3590 GFLOPs). I tempi di inferenza sono di conseguenza più rapidi (ad esempio, 13,4 ms contro 28,5 ms per VirTues).

Significato
Il documento posiziona ImmuVis come un framework pratico per la modellazione IMC nel mondo reale che affronta i vincoli critici dell'eterogeneità dei pannelli e dell'efficienza computazionale. Consentendo a un singolo modello fondazione di operare su sottoinsiemi arbitrari di marcatori senza re-addestramento, ImmuVis facilita il profilamento spaziale su larga scala dei tessuti e l'aggiornamento retrospettivo di coorti legacy. L'integrazione di un'incertezza calibrata offre un approccio consapevole dell'affidabilità alla colorazione virtuale, cruciale per le applicazioni cliniche in cui i falsi positivi o le previsioni a bassa confidenza devono essere segnalati. Gli autori suggeriscono che questa adattabilità a livello di operatore ai canali fornisce una ricetta generale per la costruzione di modelli fondazione nell'imaging multiplex e oltre, combinando flessibilità, efficienza e affidabilità all'interno di una singola architettura.