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

Deze paper introduceert RAFM, een methode die Retrieval-Augmented Flow Matching toepast om ongepaarde CBCT- naar CT-translatie in de medische beeldvorming te verbeteren door semantisch consistente pseudo-paren te genereren, wat resulteert in superieure prestaties voor synthetische CT-generatie vergeleken met bestaande methoden.

De kern

Stel je voor dat je een fotograaf bent die elke dag foto's maakt van een patiënt voor een stralingstherapie. Je hebt twee soorten camera's:

1. De 'Gouden Standaard' (CT-scan): Deze maakt prachtige, scherpe foto's met de perfecte kleuren en helderheid. Hiermee kunnen artsen precies berekenen hoeveel straling ze moeten geven.
2. De 'Werkcamera' (CBCT): Deze zit vast aan de stralingsmachine. Hij maakt foto's terwijl de patiënt op de tafel ligt, zodat de artsen precies weten waar ze moeten schieten. Maar deze foto's zijn vaak wazig, hebben rare vlekjes (artefacten) en de 'helderheidswaarden' kloppen niet. Je kunt ze dus niet direct gebruiken voor de berekening van de straling.

Het probleem:
Om de 'Werkcamera'-foto's om te zetten in 'Gouden Standaard'-foto's, zouden artsen idealiter een paar foto's nodig hebben van exact hetzelfde moment en exact dezelfde persoon. Maar dat is bijna onmogelijk. De patiënt beweegt, de organen veranderen van vorm, en de machines staan op verschillende plekken. Je hebt dus geen perfecte paren; je hebt alleen een stapel 'Werkcamera'-foto's en een stapel 'Gouden Standaard'-foto's van verschillende mensen op verschillende tijdstippen.

De oude oplossing (en waarom die faalt):
Vroeger probeerden computersystemen (zoals GANs) deze twee stapels met elkaar te vergelijken. Het was alsof je probeert een wazige foto van een man in een blauw overhemd om te zetten in een scherpe foto van een man in een rood overhemd, zonder dat je weet wie wie is. Het systeem raakt in de war, maakt rare keuzes en de resultaten zijn vaak onstabiel of onnauwkeurig.

De nieuwe oplossing: RAFM (De 'Slimme Zoeker')
De auteurs van dit paper hebben een nieuwe manier bedacht, genaamd RAFM. Ze gebruiken een slimme techniek die we "Flow Matching" noemen, maar dan met een extra trucje: Retrieval-Augmented (verrijkt met zoeken).

Hier is hoe het werkt, in simpele termen:

1. De 'Flow' (De Stroom)

Stel je voor dat je een rivier hebt. Aan de ene kant (bron) heb je de wazige CBCT-foto's. Aan de andere kant (doel) heb je de scherpe CT-foto's.
De oude methoden probeerden een brug te bouwen tussen willekeurige punten aan de bron en willekeurige punten aan de doelkant. Maar als je een wazige foto van een knie koppelt aan een scherpe foto van een hoofd, krijg je een rommelig resultaat.

Rectified Flow (de basis van hun methode) zegt: "Laten we een rechte lijn trekken tussen bron en doel." Maar in een kleine dataset (weinig patiënten) is het lastig om de juiste lijnen te vinden zonder dat je per ongeluk een knie aan een hoofd koppelt.

2. De 'Retrieval' (De Slimme Zoeker)

Hier komt de magie van RAFM om de hoek kijken. In plaats van willekeurig te zoeken, gebruikt het systeem een Slimme Zoeker (een AI die is getraind om beelden te begrijpen, genaamd DINOv3).

• De Bibliotheek: Het systeem heeft een enorme digitale bibliotheek (een 'geheugenbank') vol met scherpe CT-foto's van verschillende mensen.
• De Zoektocht: Als het systeem een wazige CBCT-foto van een knie krijgt, kijkt de Slimme Zoeker niet naar een willekeurige foto in de bibliotheek. Hij zoekt naar de foto die er het meest op lijkt qua structuur. Hij vindt een scherpe CT-foto van een knie (zelfs als die van een heel andere persoon is).
• De Koppeling: Nu koppelt het systeem de wazige knie aan de scherpe knie. Dit is een 'pseudo-paar'. Het is geen echt paar van dezelfde persoon, maar het is wel een betrouwbare match.

3. Het Resultaat: Een Gladmakende Machine

Met deze slimme koppelingen leert het systeem een rechte, stabiele lijn (een 'flow') te volgen. Het weet precies hoe het de wazigheid moet weghalen en de kleuren moet corrigeren, omdat het steeds vergelijkt met de juiste 'doel' (bijvoorbeeld: "Hoe ziet een gezonde knie eruit?").

Waarom is dit zo goed?

• Geen ruzie: Oude methoden waren als een gevecht tussen twee AI's (een maker en een criticus), wat vaak instabiel was. RAFM is als een rustige leraar die stap voor stap corrigeert.
• Geen perfecte paren nodig: Je hoeft niet dezelfde patiënt op twee verschillende momenten te scannen. Het systeem is slim genoeg om te begrijpen dat een knie van persoon A op een knie van persoon B lijkt.
• Sneller en scherper: De resultaten laten zien dat de nieuwe foto's veel scherper zijn, minder ruis hebben en de anatomie (de vorm van de organen) perfect behouden blijft.

Kort samengevat:
RAFM is als een slimme tolk die twee talen vertaalt (CBCT naar CT). In plaats van willekeurig woorden te vertalen, kijkt hij eerst in een woordenboek (de geheugenbank) om het beste synoniem te vinden dat past bij de context. Zo krijgt de patiënt een perfecte, scherpe foto voor zijn behandeling, zelfs als de oorspronkelijke foto's niet perfect waren.

Technische samenvatting

Probleemstelling

In de radiotherapie is Computertomografie (CT) de standaard voor behandelingsplanning vanwege de betrouwbare Hounsfield-eenheden (HU) die nodig zijn voor dosisberekening. Cone-beam CT (CBCT), die routinematig tijdens de behandeling wordt verkregen, is waardevol voor adaptieve radiotherapie, maar lijdt aan ernstige artefacten en onbetrouwbare HU-waarden. Dit beperkt het directe gebruik voor dosiskalkulaties.

Het genereren van synthetische CT-beelden (sCT) vanuit CBCT is daarom een cruciale taak. Het grootste obstakel is echter dat gepaarde CBCT-CT-data vaak ontbreken of onbetrouwbaar zijn door tijdsverschillen, anatomische variaties en registratiefouten. Bestaande methoden voor ongepaarde vertaling (zoals GANs of Diffusiemodellen) kampen vaak met instabiliteit in het trainingsproces, gevoeligheid voor hyperparameters of complexe trainingspiplines. Bovendien zijn medische datasets vaak klein, wat het trainen van modellen met kleine batchgroottes vereist.

Methodologie: RAFM

De auteurs introduceren Retrieval-Augmented Flow Matching (RAFM), een niet-adversariaal framework dat Gekorrigeerde Flow (Rectified Flow - RF) toepast op ongepaarde medische beeldvertaling, versterkt door een zoekstrategie.

1. Rectified Flow (RF) Basis:

   • RF modelleert de vertaling als een deterministische transportstroom tussen twee verdelingen (CBCT en CT) via een differentiaalvergelijking (ODE).
   • Het doel is om een snelheidsveld $v_\theta$ te leren dat een lineair pad volgt tussen een startpunt ($x_0$, CBCT) en een eindpunt ($x_1$, CT).
   • In theorie vereist RF geen voxel-gealigneerde paren, maar slechts een koppeling (coupling) waarbij de marginaalverdelingen overeenkomen met de bron- en doeldomeinen.

2. De Uitdaging bij Medische Data:

   • Bij kleine datasets en kleine batchgroottes leidt willekeurige of lokaal-batch-gebaseerde koppeling van CBCT- en CT-slices vaak tot semantisch mismatchende paren (bijv. een CBCT-slice van de heup wordt gekoppeld aan een CT-slice van de borstkas). Dit introduceert ruis in het trainingsdoel en schaadt de anatomische consistentie.

3. Retrieval-Augmented Strategie (RAFM):

   • Om dit op te lossen, introduceert RAFM een globale CT-geheugenbank (memory bank) die wordt gebruikt om de koppeling te verbeteren.
   • Feature Extractie: Een bevroren DINOv3-encoder wordt gebruikt om features te extraheren uit CT-slices.
   • Zoekmechanisme: Voor elke CBCT-slice in de batch wordt de meest vergelijkbare CT-slice opgezocht in de geheugenbank op basis van cosine-afstand in de feature-ruimte.
   • Pseudo-paren: Deze gevonden CT-slice vormt een "pseudo-eindpunt" ($x_1$) voor de CBCT-slice ($x_0$). Dit creëert een retrieval-gebaseerde empirische koppeling ($\rho_{retr}$) die semantisch consistenter is dan willekeurige koppeling.
   • Training: Het model wordt getraind om het snelheidsveld te leren dat deze semantisch gereserveerde paden volgt, zonder gebruik te maken van echte subject-identiteiten of gepaarde annotaties.

4. Inferentie:

   • Tijdens inferentie wordt de geleerde ODE opgelost (van $t=0$ tot $t=1$) om vanuit een CBCT-slice een synthetische CT-slice te genereren.

Belangrijkste Bijdragen

• Toepassing van RF op Medische Beeldvorming: Systematische exploratie van Rectified Flow voor ongepaarde CBCT-naar-CT vertaling, waarbij de voordelen van deterministische transport (stabiliteit, geen adversariaal trainen) worden benut.
• Retrieval-Augmented Coupling: Een innovatieve oplossing voor het probleem van kleine datasets en kleine batchgroottes. Door een globale geheugenbank en DINOv3-features te gebruiken, worden semantisch betere pseudo-paren gegenereerd, wat de kwaliteit van de transportstroom aanzienlijk verbetert.
• Strict Unpaired Protocol: Het framework werkt strikt ongepaard (geen subject-overlap tijdens training), wat dichter bij de realiteit van klinische data ligt dan veel bestaande methoden.

Resultaten

De methode is geëvalueerd op de SynthRAD2023 dataset (bekken) onder een strikt "subject-level true-unpaired" protocol.

• Kwantitatieve Prestaties: RAFM presteert consistent beter dan bestaande methoden (waaronder CycleGAN, CUT, SynDiff en UNSB) op alle metrics:
  • FID (Fréchet Inception Distance): 53.29 (laagste, wat betekent beste distributierealiteit).
  • MAE (Mean Absolute Error): 101.2 HU (laagste fout).
  • SSIM & PSNR: Hoogste scores voor structurele gelijkenis en beeldkwaliteit.
  • SegScore: 75.77% (hoogste anatomische consistentie, gemeten via segmentatie van organen).
• Kwalitatieve Resultaten: Visuele vergelijkingen tonen minder artefacten en stabielere anatomische structuren in vergelijking met concurrenten.
• Efficiëntie: RAFM is sneller in inferentie dan diffusiemodellen (alleen 10 ODE-stappen nodig) en vereist minder rekenkracht dan bidirectionele frameworks zoals CycleGAN.

Betekenis en Conclusie

RAFM bewijst dat Rectified Flow een krachtig alternatief is voor GANs en Diffusiemodellen in de medische beeldanalyse, mits het probleem van slechte koppeling bij kleine datasets wordt opgelost. De introductie van retrieval-verrijkte koppeling is een sleutelfactor die de kloof tussen ongepaarde en gepaarde training verkleint.

De studie benadrukt dat de kwaliteit van de koppeling (hoe goed de start- en eindpunten semantisch overeenkomen) cruciaal is voor het behoud van anatomische structuur in ongepaarde vertaling. RAFM biedt een stabiel, niet-adversariaal en efficiënt kader voor het genereren van hoogwaardige synthetische CT-beelden, wat direct bijdraagt aan de verbetering van adaptieve radiotherapie workflows.