RAFM: Retrieval-Augmented Flow Matching for Unpaired CBCT-to-CT Translation

Die vorgestellte Arbeit stellt RAFM vor, eine Methode zur unpaarigen CBCT-zu-CT-Übersetzung in der medizinischen Bildgebung, die durch den Einsatz von Retrieval-verstärktem Flow Matching und einem globalen CT-Speicherbank die Stabilität des Trainings verbessert und damit die Qualität synthetischer CT-Bilder für die Strahlentherapie signifikant steigert.

Das Wesentliche

Stell dir vor, du bist ein Architekt, der einen perfekten Bauplan für ein Haus (einen CT-Scan) braucht, um die genaue Menge an Ziegelsteinen (Strahlendosis) zu berechnen, die für eine Renovierung (Strahlentherapie) nötig ist.

Das Problem: Die Architekten haben oft nur alte, verwaschene Fotos vom Haus (CBCT-Scans). Diese Fotos sind zwar schnell gemacht, aber sie sind unscharf, haben seltsame Schatten (Artefakte) und die Farben sind so verfälscht, dass man die Ziegelsteine nicht zählen kann. Wenn man diese Fotos direkt nutzt, könnte die Renovierung katastrophal enden.

Bisherige Methoden, um aus dem unscharfen Foto einen perfekten Bauplan zu machen, waren wie ein wildes Ratespiel:

1. GANs (Generative Adversarial Networks): Zwei KI-Modelle, die gegeneinander kämpfen (wie ein Fälscher und ein Detektiv). Das funktioniert oft gut, ist aber instabil und das Ergebnis kann manchmal "halluzinieren" (falsche Wände erfinden).
2. Diffusionsmodelle: Ein sehr langsamer Prozess, bei dem man aus einem Bild mit viel Rauschen nach und nach ein klares Bild entwickelt. Das ist rechenintensiv und langsam.

Die neue Lösung: RAFM (Der "Erinnerungs-Transporter")

Die Autoren dieses Papers haben eine neue Methode namens RAFM entwickelt. Sie nutzen eine Technik namens "Flow Matching" (Fließgleichgewicht), die man sich wie eine gerade, geradlinige Autobahn zwischen zwei Städten vorstellen kann.

Das Problem mit der Autobahn:
Normalerweise versucht die KI, jedes CBCT-Foto mit einem zufälligen CT-Bauplan zu verbinden, um die "Abfahrt" (den Weg) zu lernen. Aber da die Datenbanken klein sind und die Fotos oft sehr unterschiedlich aussehen (z. B. ein Foto von einem Patienten mit einem Beinbruch und ein Bauplan von einem gesunden Patienten), passt die Verbindung nicht. Es ist, als würde man versuchen, ein Foto von einem Wald mit einem Bauplan für eine Wüste zu verbinden. Das Ergebnis ist chaotisch.

Die geniale Idee: Der "Gedächtnis-Speicher"
RAFM löst dieses Problem mit einem cleveren Trick, den sie Retrieval-Augmented (erweiterter Abruf) nennen:

1. Die Bibliothek: Die KI hat eine riesige, vorgefertigte Bibliothek mit tausenden perfekten CT-Bauplänen (dem "CT-Memory Bank").
2. Der Sucher (DINOv3): Bevor die KI den Weg berechnet, schaut sie sich das unscharfe CBCT-Foto an und fragt einen sehr klugen Sucher (einen eingefrorenen KI-Encoder namens DINOv3): "Hey, welcher perfekte Bauplan in unserer Bibliothek sieht diesem unscharfen Foto am ähnlichsten?"
3. Die perfekte Paarung: Statt einen zufälligen Plan zu nehmen, holt sich die KI den besten passenden Plan aus der Bibliothek.
4. Der Transport: Jetzt baut die KI eine stabile, gerade Straße zwischen dem unscharfen Foto und diesem passenden Plan. Sie lernt, wie man das eine in das andere verwandelt, ohne die Struktur des Hauses (die Anatomie des Patienten) zu zerstören.

Warum ist das so toll?

• Kein "Paar"-Problem: Normalerweise braucht man für solche Aufgaben ein Foto und den exakten dazugehörigen Plan vom gleichen Patienten zur gleichen Zeit. Das ist in der Medizin oft unmöglich, weil Patienten sich bewegen oder die Scans zu unterschiedlichen Zeiten gemacht wurden. RAFM kommt ohne diese perfekten Paare aus, indem es einfach die ähnlichsten Paare aus der großen Bibliothek sucht.
• Stabilität: Weil die Verbindung zwischen Foto und Plan logisch ist (ähnliche Anatomie), lernt die KI viel schneller und macht weniger Fehler.
• Geschwindigkeit: Im Gegensatz zu den langsamen Diffusionsmodellen ist RAFM schnell wie ein Sprinter auf der Autobahn.

Das Ergebnis im Test

Die Forscher haben ihre Methode an einem großen Datensatz (SynthRAD2023) getestet, bei dem sie die Trainingsdaten so streng getrennt haben, dass keine einzige Information über die Patientenpaarung durchkam (ein "echtes" unpaariges Szenario).

Das Ergebnis: RAFM war der Gewinner.

• Die Bilder waren schärfer (weniger Fehler in den Zahlenwerten).
• Die Anatomie (Knochen, Organe) wurde besser erhalten als bei allen anderen Methoden.
• Die KI "halluzinierte" weniger falsche Strukturen.

Zusammenfassend:
Stell dir RAFM wie einen hochintelligenten Übersetzer vor, der nicht einfach zufällige Wörter austauscht, sondern erst im Wörterbuch nachschaut, welches Wort in der Ziel-Sprache (CT) am besten zum Kontext (CBCT) passt. So entsteht eine perfekte Übersetzung, auch wenn man die Originaltexte nie direkt nebeneinander hatte. Das macht die Strahlentherapie sicherer und präziser.

Technische Zusammenfassung

1. Problemstellung

In der Strahlentherapie ist die Computertomographie (CT) der Standard für die Behandlungsplanung, da ihre Hounsfield-Einheiten (HU) zuverlässige Informationen zur Elektronendichte für die Dosisberechnung liefern. Die konventionelle Bildgebung mittels Cone-Beam-CT (CBCT) wird zwar routinemäßig zur Bildführung während der Behandlung eingesetzt, leidet jedoch unter schweren Artefakten und unzuverlässigen HU-Werten, was eine direkte Nutzung für die Dosisberechnung verhindert.

Das Ziel ist daher die Synthese von CT-ähnlichen Bildern (sCT) aus CBCT-Daten. Ein zentrales Hindernis ist jedoch die Verfügbarkeit von Daten:

• Fehlende Paarungen: Geparre CBCT-CT-Datensätze sind oft schwer zu erstellen oder aufgrund von zeitlichen Lücken, anatomischen Variationen und Registrierungsfehlern unzuverlässig.
• Ungepaartes Lernen: Daher muss das Problem oft im „ungepaarten" (unpaired) Setting gelöst werden, bei dem keine voxelgenaue Korrespondenz zwischen den CBCT- und CT-Bildern vorliegt.
• Herausforderungen bei Flow Matching: Während die Methode des „Rectified Flow" (RF) theoretisch für ungepaartes Lernen geeignet ist (da sie nur Verteilungen koppelt, keine Pixel-Paare), scheitert die praktische Anwendung in der medizinischen Bildgebung oft an kleinen Datensätzen und kleinen Batch-Größen. Zufällige oder nur batch-lokale Pseudo-Paarungen führen zu semantisch inkonsistenten Endpunkten und instabilen Transportzielen, was die anatomische Integrität gefährdet.

2. Methodik: Retrieval-Augmented Flow Matching (RAFM)

Die Autoren schlagen RAFM vor, eine Erweiterung des Rectified Flow (RF), die durch einen retrieval-basierten Mechanismus die Qualität der Pseudo-Paarung verbessert, ohne dabei auf gepaarte Daten zurückzugreifen.

Kernkomponenten:

• Rectified Flow (RF): Das Modell lernt ein deterministisches Vektorfeld $v_\theta(x, t)$, das eine gerade Linie (ODE) zwischen der Quellverteilung (CBCT) und der Zielverteilung (CT) beschreibt. Das Trainingsziel ist die Regression auf die konstante Geschwindigkeit $x_1 - x_0$ entlang dieses Pfades.
• Globaler CT-Speicher (Memory Bank): Um das Problem der schlechten Paarung bei kleinen Batches zu lösen, wird ein globaler Speicher für CT-Bilder aufgebaut.
  • Ein eingefrorener DINOv3-Encoder extrahiert Merkmale aus jedem CT-Slice.
  • Diese Merkmale werden in einer FIFO-Warteschlange (Memory Bank) gespeichert, die über die Trainingsiterationen hinweg aktualisiert wird (Rolling Buffer).
• Retrieval-Guided Coupling: Für jeden CBCT-Slice im aktuellen Batch wird der ähnlichste CT-Slice aus dem globalen Speicher basierend auf der Kosinus-Ähnlichkeit der DINOv3-Merkmale gesucht.
  • Dies bildet ein „Pseudo-Paar" $(x_0, x_1)$, das semantisch konsistenter ist als eine zufällige Paarung.
  • Wichtig: Es werden keine Patient-IDs oder zeitlichen Korrespondenzen verwendet; das System bleibt strikt ungepaart.
• Architektur: Das Vektorfeld wird durch ein zeit-konditioniertes U-Net parametrisiert. Beim Inferenz wird die gelernte ODE von $t=0$ (CBCT) bis $t=1$ (synthetisches CT) gelöst.

3. Schlüsselbeiträge

1. Einführung von Retrieval in Flow Matching: RAFM ist das erste Framework, das Retrieval-Augmentation in den Kontext des ungepaarten medizinischen Bild-zu-Bild-Transfers mittels Rectified Flow integriert.
2. Lösung des Small-Data-Problems: Durch die Nutzung eines globalen Merkmals-Speichers wird die Limitierung kleiner Batch-Größen umgangen, was zu einer stabileren und semantisch korrekteren Kopplung (Coupling) führt.
3. Striktes Unpaired-Protokoll: Die Methode funktioniert effektiv unter einem strengen „Subject-Level True-Unpaired"-Protokoll, bei dem im Training keine einzige Patientenkombination aus CBCT und CT vorhanden ist.
4. Non-Adversarialer Ansatz: Im Gegensatz zu GAN-basierten Methoden bietet RAFM eine stabile, nicht-adversative Optimierung, die für die Erhaltung der Anatomie in der Medizin entscheidend ist.

4. Ergebnisse

Die Methode wurde auf dem SynthRAD2023-Datensatz (Beckenbereich) evaluiert und mit State-of-the-Art-Methoden (GANs wie CycleGAN, CUT sowie Diffusion/SB-Modelle wie SynDiff, UNSB) verglichen.

• Quantitative Leistung: RAFM übertraf alle Vergleichsmethoden in allen Metriken:
  • MAE (Mean Absolute Error): 101.2 HU (niedriger als 104.2 bei SynDiff und 110.8 bei UNSB).
  • FID (Fréchet Inception Distance): 53.29 (deutlich besser als 62.91 bei UNSB), was eine höhere Verteilungsrealität anzeigt.
  • SSIM & PSNR: Höhere Werte für strukturelle Ähnlichkeit und Bildqualität.
  • SegScore: 75.77% (basierend auf der Segmentierung von Organen mit TotalSegmentator), was die hervorragende Erhaltung der anatomischen Struktur beweist.
• Qualitative Ergebnisse: Visuelle Vergleiche zeigen eine sauberere Unterdrückung von Artefakten und stabilere anatomische Strukturen im Vergleich zu GAN- und Diffusions-basierten Ansätzen.
• Effizienz: RAFM ist beim Inferenz deutlich schneller als Multi-Step-Diffusionsmodelle (nur 10 ODE-Schritte) und benötigt weniger Rechenressourcen als bidirektionale GAN-Frameworks.

Ablationsstudie:
Die Studie zeigte, dass die Qualität der Kopplung der entscheidende Faktor ist. Während zufällige Paarungen (Random Coupling) bereits besser sind als reine Regression (Vanilla U-Net), führt die Erweiterung des Kandidatenpools durch das Retrieval (Memory Bank) zu signifikanten Verbesserungen. Ein Speicher von $K=512$ erwies sich als optimaler Kompromiss.

5. Bedeutung und Fazit

RAFM adressiert eine kritische Lücke in der medizinischen Bildverarbeitung: die Generierung hochwertiger synthetischer CT-Bilder aus CBCT-Daten ohne Zugang zu gepaarten Trainingsdaten.

• Praktische Relevanz: Da gepaarte Daten in der klinischen Praxis oft nicht verfügbar sind, ermöglicht RAFM eine robuste Anpassung an adaptive Strahlentherapie-Workflows.
• Methodischer Fortschritt: Die Arbeit demonstriert, dass die Kombination von modernen Vision-Transformern (DINOv3) für das Retrieval mit Flow-Matching-Modellen ein vielversprechender Weg ist, um die Stabilität und Qualität von ungepaarten Übersetzungsmodellen in datenarmen Szenarien zu verbessern.
• Zukunftsperspektive: Der Code ist öffentlich verfügbar, was die Reproduzierbarkeit und weitere Forschung in diesem Bereich fördert.

Zusammenfassend stellt RAFM einen neuen State-of-the-Art für die ungepaarte CBCT-zu-CT-Übersetzung dar, der durch eine intelligente, merkmalsbasierte Paarungsstrategie die Nachteile kleiner medizinischer Datensätze kompensiert.