MUST: Modality-Specific Representation-Aware Transformer for Diffusion-Enhanced Survival Prediction with Missing Modality

1. Probleemstelling

In de precisiestofologie is het combineren van multimodale medische data (zoals Whole Slide Images (WSI) van weefsel en genomische profielen) essentieel voor nauwkeurige overlevingsvoorspelling. Echter, een groot obstakel voor de klinische toepassing is dat data vaak onvolledig is. Door kosten, technische beperkingen of retrospectieve datacollecties ontbreken vaak bepaalde modaliteiten (bijv. wel histopathologie maar geen genoomdata, of vice versa).

Bestaande methoden voor het omgaan met ontbrekende modaliteiten hebben fundamentele beperkingen:

• Feature Alignment: Deze methoden dwingen representaties gelijk te maken zonder expliciet te modelleren welke informatie specifiek is voor een modaliteit en welke gedeeld is.
• Imputatie: Generatieve modellen proberen ontbrekende features te synthetiseren, maar werken vaak in hoge dimensies wat leidt tot ruis en degradeert de prestaties.
• Gedeelde verdeling: Methoden zoals LD-CVAE leren gezamenlijke posteriors, maar kunnen niet duidelijk onderscheid maken tussen wat uniek is voor een modaliteit en wat uit andere modaliteiten kan worden afgeleid.

Het kernprobleem is het gebrek aan expliciete modellering van modality-specific versus cross-modally inferable informatie.

2. Methodologie: MUST Framework

De auteurs stellen MUST (Modality-Specific representation-aware Transformer) voor, een raamwerk dat expliciet de representaties van elke modaliteit decomposeert in twee componenten:

1. Modality-specific componenten: Informatie die uniek is voor die modaliteit en niet uit andere kan worden afgeleid.
2. Gedeelde (cross-modal) componenten: Informatie die gemeenschappelijk is en uit elke beschikbare modaliteit kan worden afgeleid.

Het proces verloopt in drie hoofdfasen:

A. Algebraïsche Decompositie in een Low-Rank Subruimte

Om algebraïsche operaties mogelijk te maken tussen verschillende semantische ruimtes (weefsel vs. genoom), projecteren de auteurs de representaties naar een leerbaar low-rank gedeelde subruimte.

• Bidirectionele Cross-Attention: Er wordt interactie tussen modaliteiten gemodelleerd om contextuele componenten ($c_{P \leftarrow G}$ en $c_{G \leftarrow P}$) te extraheren.
• Algebraïsche Constraints: De globale representatie ($g$) wordt ontbonden in een uniek deel ($\hat{u}$) en een gedeeld deel ($\hat{c}$) via projectie-operatoren.
  • $g_P = \hat{u}_P + \hat{c}_{G \leftarrow P}$
  • $g_G = \hat{u}_G + \hat{c}_{P \leftarrow G}$
• Orthogonaliteit: Er worden strakke constraints opgelegd om te garanderen dat de unieke componenten orthogonaal zijn aan elkaar en aan het gedeelde deel. Dit zorgt ervoor dat het gedeelde deel deterministisch uit elke beschikbare modaliteit kan worden afgeleid.

B. Progressieve Trainingsstrategie

Om stabiele convergentie te bereiken, wordt een twee-staps training gebruikt:

1. Stage 1 (Warm-up): De encoders leren taakrelevante signalen met toegevoegde Gaussische ruis om overfitting op schijnbare correlaties te voorkomen.
2. Stage 2 (Decompositie): Het volledige raamwerk wordt getraind met een verliesfunctie die survival-predicatie combineert met algebraïsche consistentie, gedeelde consistentie en orthogonaliteitsverliezen.

C. Latent Diffusion Models (LDM) voor Ontbrekende Modaliteiten

Voor de informatie die echt uniek is voor een ontbrekende modaliteit (die niet deterministisch kan worden gereconstrueerd), wordt een Conditionele Latent Diffusion Model gebruikt.

• Wanneer een modaliteit ontbreekt (bijv. genoom), wordt het unieke deel ($\hat{u}_G$) gegenereerd via DDIM-sampling.
• De generatie is geconditioneerd op:
  1. Het gereconstrueerde gedeelde component ($\hat{c}$), afgeleid van de aanwezige modaliteit (deterministisch).
  2. Een structurele prior (een leerbare [CLS]-token) die de structuur van de ontbrekende modaliteit encodeert.
• Dit isoleert de stochastische generatie tot alleen de unieke residuen, wat de kwaliteit en betrouwbaarheid verhoogt.

3. Belangrijkste Bijdragen

1. Expliciete Algebraïsche Decompositie: Een architectuur die via bidirectionele cross-attention en algebraïsche constraints in een low-rank subruimte precies identificeert welke informatie verloren gaat bij een ontbrekende modaliteit. Dit maakt deterministische herwinning van het gedeelde deel mogelijk.
2. Twee-staps Reconstructiestrategie: Een hybride aanpak die deterministische algebraïsche projectie combineert voor gedeelde informatie en probabilistische generatie (via LDM) voor modality-specifieke informatie. Dit beperkt stochastische onzekerheid tot alleen de unieke residuen.
3. State-of-the-Art Prestaties: MUST presteert superieur op vijf TCGA-kankerdatasets, zowel bij volledige data als bij scenario's met ontbrekende modaliteiten, met klinisch acceptabele inferentielatentie.

4. Resultaten

De methoden zijn getest op vijf TCGA-datasets (BLCA, BRCA, GBMLGG, LUAD, UCEC) met als maatstaf de C-index (concordance index).

• Volledige Data: MUST bereikt een C-index van 0.742, wat een verbetering is van 2.5% ten opzichte van de huidige state-of-the-art (CMTA, 0.724).
• Ontbrekende Genomica: MUST behoudt een C-index van 0.716 (slechts 3.5% daling), wat significant beter is dan concurrenten zoals LD-CVAE (0.697) en ShaSpec (0.676).
• Ontbrekende Pathologie: MUST behoudt een C-index van 0.739 (slechts 0.4% daling). Opmerkelijk is dat MUST op sommige datasets zelfs beter presteert bij ontbrekende pathologie dan bij volledige data, waarschijnlijk omdat het diffusiemodel ruis in de pathologische patches filtert.
• Klinische Validatie: Kaplan-Meier-curves tonen aan dat de risicogroepering (hoog vs. laag risico) statistisch significant blijft (p < 0.05) zelfs bij ontbrekende modaliteiten, wat essentieel is voor klinische besluitvorming.

5. Betekenis en Conclusie

MUST lost een fundamenteel probleem op in de multimodale medische AI: de onzekerheid bij ontbrekende data. Door expliciet te scheiden tussen wat gedeeld is (deterministisch herwinbaar) en wat uniek is (probabilistisch te genereren), vermijdt het de valkuilen van eerdere methoden die ofwel te veel ruis introduceren ofwel de unieke informatie negeren.

De algebraïsche invertibiliteit is de kerninnovatie: het gedeelde deel is altijd afleidbaar uit de aanwezige data, wat een betrouwbare basis biedt voor de generatie van het ontbrekende unieke deel. Dit maakt MUST zeer geschikt voor de klinische praktijk, waar data vaak onvolledig is, en biedt een robuust alternatief voor methoden die volledige data vereisen. De auteurs merken ook op dat het raamwerk potentieel heeft voor training met onvolledige paren, hoewel dit nog verder onderzoek vereist.