BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability

Technischer Zusammenfassung von "BrainAnytime: Anatomiebewusstes Cross-Modal Pretraining für die Analyse von Gehirnimages mit beliebiger Modalitätsverfügbarkeit"

Problemstellung

Klinische Diagnoseworkflows für neurodegenerative Erkrankungen, insbesondere die Alzheimer-Krankheit (AD), folgen einem „Pfad der Modalitätssteigerung". Kliniker beginnen typischerweise mit routinemäßigen strukturellen Bildgebungsverfahren (z. B. T1-gewichtete MRT), fügen selektiv Sequenzen hinzu (z. B. FLAIR, T2), um Differentialdiagnosen zu verfeinern, und reservieren molekulare Bildgebung (z. B. Amyloid-PET) für unklare Fälle. Folglich sind Patientendaten in realen Kohorten durch heterogene und oft unvollständige Modalitäts-Teilmengen gekennzeichnet.

Aktuelle KI-Modelle gehen im Allgemeinen von festen Eingangsmodalitäten aus und erfordern einen Wechsel der Architektur oder ein Nachtrainieren, wenn sich die Datenverfügbarkeit ändert. Darüber hinaus leiden bestehende groß angelegte Pretraining-Modelle für Gehirnimages in diesem Kontext oft unter drei Einschränkungen:

1. Sie beschränken sich typischerweise auf MRT, wodurch sie keine expliziten Brücken zu Amyloid-PET, einem kritischen Biomarker, lernen.
2. Sie gehen von festen Modalitätskombinationen aus, was der stufenabhängigen Unvollständigkeit realer Daten nicht entspricht.
3. Das Standard-Masked-Autoencoder (MAE)-Pretraining beruht auf einer räumlich uniformen Maskierung und ignoriert die Tatsache, dass AD eine selektive regionale Vulnerabilität (z. B. im Hippocampus) aufweist und keine uniforme Pathologie.

Methodik

Die Autoren schlagen BrainAnytime vor, ein einheitliches Pretraining-Framework, das die Analyse von Gehirnimages unter beliebiger Modalitätsverfügbarkeit unterstützt. Das Framework wurde auf 34.899 3D-Gehirnscans aus fünf Datensätzen (A4, DOD-ADNI, AIBL, BraTS23, NACC) vortrainiert, die Multi-Sequence-MRT und Amyloid-PET umfassen. Es integriert drei Kernkomponenten innerhalb einer gemeinsamen 3D-Masked-Autoencoder-Architektur:

1. Multi-MAE3D (3D Multi-Modal Masked Autoencoder)

• Architektur: Ein gemeinsamer Transformer-Encoder verarbeitet jede Teilmenge verfügbarer Modalitäten (T1, T2, FLAIR, PET).
• Tokenisierung: Jede Modalität verwendet eine modalitätsspezifische 3D-Convolutional-Patch-Einbettung.
• Umgang mit fehlenden Daten: Fehlende Modalitäten werden vollständig maskiert. Um eine konstante Token-Anzahl unabhängig von der Eingangsverfügbarkeit zu gewährleisten, wird das sichtbare Token-Budget unter den beobachteten Modalitäten mithilfe einer Dirichlet-Verteilung (Dir($\alpha$)) neu verteilt.
• Decodierung: Jede Modalität verfügt über einen leichten Decoder, der alle Patches rekonstruiert. Der Trainingsverlust ($L_{MAE}$) wird nur auf maskierten Patches unter Verwendung des mittleren quadratischen Fehlers mit einer Patch-spezifischen Normalisierung berechnet.

2. Reziproke Cross-Modal Distillation (RCMD)

• Ziel: Das Erlernen von Korrespondenzen zwischen strukturellem MRT und molekularem PET.
• Mechanismus: Ein Exponential Moving Average (EMA)-Lehrernetzwerk teilt sich die Encoder-Architektur. Für gepaarte MRT-PET-Proben werden Repräsentationen auf Gruppenebene ($z_{MRT}$ und $z_{PET}$) durch Average Pooling berechnet.
• Verlust: Zwei symmetrische MLP-Prädiktoren ($g_{M\to P}$ und $g_{P\to M}$) werden trainiert, um cross-modale Lehrermerkmale vorherzusagen. Der Verlust ($L_{RCMD}$) maximiert die Kosinus-Ähnlichkeit zwischen den vom Schüler vorhergesagten Merkmalen und den Merkmalen des Lehrers im Raum der gegenüberliegenden Modalität.

3. Pathologiebewusstes Curriculum-Masking (PACM)

• Ziel: Priorisierung der krankheitsanfälligen Anatomie während der Rekonstruktion anstelle einer uniformen Maskierung.
• Bewertung der Wichtigkeit: Eine statische Bewertung wird aus neuropathologischen Prioritäten (Braak-Staging) abgeleitet, die über den AAL116-Atlas zugeordnet werden. Eine dynamische Bewertung wird durch Aggregation von Attention-Karten eines EMA-Anatomie-Lehrers erhalten. Diese werden kombiniert, um Patch-Wichtigkeitswerte zu generieren.
• Curriculum-Planung: Eine dreiphasige Strategie wird angewendet:
  1. Frühe Phase: Uniforme zufällige Maskierung.
  2. Mittlere Phase: Der Temperaturparameter ($\tau$) wird durch kosinusbasiertes Abklingen (cosine decay) angepasst, wodurch die Maskenverteilung schrittweise verfeinert wird, um Masken in Bereichen mit hoher Wichtigkeit zu bevorzugen (d. h. das Modell wird gezwungen, kritische Regionen zu rekonstruieren).
  3. Späte Phase: Feste Temperatur.
• Auswahl: Patches werden unter Verwendung des Gumbel-top-k-Tricks ausgewählt, wodurch Patches mit hoher Wichtigkeit mit größerer Wahrscheinlichkeit maskiert und somit rekonstruiert werden.

Training und Feinabstimmung:
Der Gesamtverlust beträgt $L = L_{MAE} + \lambda \cdot \mathbb{1}_{[paired]} \cdot L_{RCMD}$. Für nachgelagerte Aufgaben wird nur der Schüler-Encoder mit einem leichten Aufgaben-Head verwendet. Fehlende Modalitäten bleiben während der Inferenz auf Null gesetzt und werden durch Attention-Blöcke blockiert, wodurch sichergestellt wird, dass das Modell beliebige Teilmengen ohne Architekturänderungen verarbeitet.

Hauptbeiträge

1. Einheitliches Framework: Ein einzelnes Modell, das jede Teilmenge von Modalitäten akzeptieren kann (von einem einzelnen T1-Scan bis hin zu einer vollständigen multimodalen Untersuchung) ohne Nachtraining oder Architekturwechsel.
2. Cross-Modal-Lernen: Lernt explizit strukturell-molekulare Korrespondenzen zwischen MRT und PET durch reziproke Distillation und schließt die Lücke zwischen strukturellen und pathologischen Befunden.
3. Anatomiebewusstes Pretraining: Führt eine Curriculum-Masking-Strategie ein, die die rekonstruktive Leistung des Modells auf AD-relevante Neuroanatomie richtet und die Einschränkungen der uniformen Maskierung in Standard-MAEs adressiert.
4. Skalierbares Pretraining: Erfolgreiches Vortraining auf einem groß angelegten, heterogenen Datensatz (34.899 Scans), der klinische Muster der Datenunvollständigkeit nachahmt.

Experimentelle Ergebnisse

Das Modell wurde auf vier nachgelagerten Aufgaben (CN vs. AD, CN vs. MCI, MMSE-Regression, Altersvorhersage) über fünf klinisch motivierte Modalitätseinstellungen (von T1-only bis T1+T2+FLAIR+PET) evaluiert.

• Leistung: BrainAnytime übertraf modalitätsspezifische Modelle (z. B. 3D ResNet, M3T), Dual-Modalitäts-Fusionsmethoden, groß angelegte Gehirn-Foundation-Modelle (BrainIAC, SAM-Brain3D) und unvollständige multimodale Lernansätze (Flex-MoE, FuseMoE).
• Schlüsselmetriken:
  • CN vs. AD: Erreichte eine durchschnittliche Genauigkeit von 89,4 % (vs. 84,2 % für die nächstbeste Baseline, Flex-MoE), was einer relativen Verbesserung von 6,2 % gegenüber der stärksten Baseline mit fehlender Modalität entspricht.
  • CN vs. MCI: Erreichte eine durchschnittliche Genauigkeit von 67,2 % (vs. 60,6 % für MoE-retriever), eine relative Verbesserung von 7,0 %.
  • Regression: Zeigte überlegene Leistung bei der MMSE- und Altersvorhersage über alle Modalitätseinstellungen hinweg.
• Robustheit: Eine innerhalb von Probanden durchgeführte Robustheitsstudie bestätigte, dass sich die Leistung konsistent verbesserte, je mehr Modalitäten hinzugefügt wurden (z. B. stieg die F1 für CN vs. AD von 65,0 % bei T1-only auf 80,4 % bei voller Eingabe), was die Fähigkeit des Modells demonstriert, verfügbare Daten effektiv zu nutzen.
• Ablationsstudie: Das individuelle Entfernen von RCMD oder PACM führte zu Leistungseinbußen, was die Notwendigkeit sowohl der cross-modalen Distillation als auch der anatomiebewussten Maskierung bestätigt.

Bedeutung

Die Arbeit behauptet, dass BrainAnytime eine kritische Lücke zwischen der klinischen Realität (gestaffelte, unvollständige Datenerfassung) und aktuellen KI-Paradigmen (Modelle mit festem Input) schließt. Indem es einem einzelnen Modell ermöglicht, flexibel jede verfügbare Bildgebung zu akzeptieren, passt es sich dem klinischen Pfad der Modalitätssteigerung an. Das Framework zeigt, dass Pretraining mit cross-modaler Distillation und pathologiebewusster Maskierung ein robustes Foundation-Modell hervorbringt, das bestehende Baselines signifikant übertrifft, insbesondere in Szenarien, in denen Daten unvollständig sind oder Modalitäten fehlen. Dieser Ansatz bietet eine praktische Lösung für den Einsatz von KI in realen klinischen Umgebungen, in denen vollständige multimodale Daten selten für jeden Patienten verfügbar sind.