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

Cet article présente RAFM, une méthode d'apprentissage non apparié pour la traduction CBCT-CT qui améliore la stabilité et la qualité de la génération d'images synthétiques en intégrant un mécanisme de récupération guidé par DINOv3 dans le cadre du Flow Matching rectifié.

L'essentiel

🏥 Le Problème : La photo floue et la carte précise

Imaginez que vous êtes un médecin radiothérapeute. Pour soigner un cancer, vous avez besoin de deux choses :

1. Une photo en temps réel de la tumeur pour viser juste pendant la séance. C'est ce qu'on appelle le CBCT (Cone-Beam CT). C'est comme une photo prise avec un appareil photo basique : on voit la forme, mais l'image est pleine de "bruit" (artefacts), floue, et les couleurs (les niveaux de densité) ne sont pas fiables.
2. Une carte de précision absolue pour calculer la dose de radiation exacte. C'est le CT classique. C'est comme une carte topographique parfaite : tout est net, précis, et les valeurs sont fiables.

Le dilemme : On ne peut pas toujours avoir la "carte parfaite" (CT) au moment où on prend la "photo floue" (CBCT). Parfois, le patient a bougé, ou les machines ne sont pas synchronisées. On se retrouve donc avec une photo floue qu'on aimerait transformer en carte précise, mais sans avoir le modèle original (la vraie carte) juste à côté pour faire la comparaison.

C'est comme essayer de restaurer une vieille photo abîmée sans avoir la photo originale sous les yeux.

🎨 La Solution : RAFM (Le traducteur intelligent)

Les chercheurs ont créé une nouvelle méthode appelée RAFM (Retrieval-Augmented Flow Matching). Pour comprendre comment ça marche, utilisons une analogie culinaire.

1. L'ancien problème : Le chef qui devine au hasard

Avant, les ordinateurs essayaient de transformer l'image floue en image nette en faisant des paris au hasard.

• Imaginez un chef qui doit transformer un plat de pâtes (l'image floue) en un plat de sushi (l'image nette).
• S'il prend une photo de n'importe quel plat de pâtes et la compare à n'importe quelle photo de sushi au hasard, il va se tromper. Il pourrait essayer de transformer des spaghettis en sashimi de saumon, ce qui est absurde. Le résultat sera une soupe bizarre.
• En mathématiques, on appelle cela un "appariement aléatoire". Avec peu de données (comme en médecine), ce hasard crée trop d'erreurs.

2. La nouvelle astuce : La bibliothèque de référence (Le "Mémoire Bank")

RAFM change la donne en ajoutant un bibliothécaire intelligent.

Voici comment le processus fonctionne, étape par étape :

• L'Encyclopédie Gelée (DINOv3) : Avant même de commencer, l'ordinateur a lu des milliers de photos de "plats parfaits" (des images CT saines) et a mémorisé leur "goût" (leurs caractéristiques visuelles) dans une grande bibliothèque virtuelle. Il ne les a pas mémorisées par nom, mais par leur "saveur" visuelle.
• La Recherche (Retrieval) : Quand le système reçoit une image floue (CBCT) d'un patient, il ne devine pas au hasard. Il va dans sa bibliothèque et dit : "Attends, cette image de pâtes a une saveur très proche de ce plat de sushi spécifique que j'ai en mémoire."
• Le Couple Idéal (Pseudo-pair) : Au lieu de comparer l'image floue à n'importe quelle image nette, il la compare à l'image nette la plus similaire trouvée dans sa bibliothèque. C'est comme si le chef trouvait le plat de pâtes qui ressemble le plus au sushi qu'il veut créer.
• Le Chemin Droit (Flow Matching) : Une fois qu'il a trouvé ce "cousin" parfait, l'ordinateur trace une ligne droite imaginaire entre l'image floue et l'image nette. Il apprend à glisser doucement le long de cette ligne pour transformer le flou en net, sans faire de sauts brusques.

🚀 Pourquoi c'est génial ?

1. Pas de combat (Non-adversarial) : Les anciennes méthodes utilisaient deux IA qui se battaient l'une contre l'autre (comme un faussaire et un détective), ce qui rendait l'apprentissage instable et difficile. RAFM, lui, suit une ligne droite logique. C'est plus calme et plus stable.
2. Idéal pour les petits groupes : En médecine, on a souvent peu de patients. Les méthodes précédentes échouaient quand il y avait peu de données. RAFM, grâce à sa "bibliothèque" (Memory Bank), peut trouver des ressemblances même avec un petit nombre d'exemples.
3. Résultats époustouflants : Dans les tests, RAFM a produit des images beaucoup plus nettes, avec moins d'erreurs de calcul, et a mieux préservé la forme des organes (comme le bassin) que toutes les méthodes précédentes.

🏁 En résumé

Imaginez que vous devez traduire un texte écrit à la main, illisible, en un texte dactylographié parfait, mais vous n'avez pas le texte original sous les yeux.

• Les anciennes méthodes prenaient un mot au hasard dans un dictionnaire et espéraient que ça colle. Ça faisait des phrases sans sens.
• RAFM, lui, consulte une immense bibliothèque de textes similaires. Il trouve le mot qui correspond le mieux au contexte, puis il recopie le texte en suivant une ligne droite logique.

Grâce à cette astuce, les médecins pourront bientôt utiliser les images floues prises directement sur la machine de radiothérapie pour calculer des traitements ultra-précis, sans avoir besoin d'attendre une autre machine ou de faire des examens supplémentaires. C'est un pas de géant pour la sécurité des patients !

Résumé technique

1. Problématique

La tomographie par ordinateur à cône (CBCT) est couramment utilisée en radiothérapie pour le guidage des traitements, mais elle souffre de deux limitations majeures :

• Artéfacts sévères et valeurs d'Unité Hounsfield (HU) peu fiables, ce qui empêche son utilisation directe pour le calcul de dose.
• Difficulté d'appariement des données : La génération d'images CT synthétiques (sCT) à partir de CBCT nécessite idéalement des données appariées (CBCT-CT). Cependant, en pratique, ces paires sont souvent indisponibles ou peu fiables en raison d'écarts temporels, de variations anatomiques et d'erreurs de recalage.

Les méthodes existantes (GANs, modèles de diffusion) pour la traduction non appariée souffrent souvent d'instabilité à l'entraînement, de sensibilité aux hyperparamètres ou de pipelines complexes. L'objectif est donc de développer une méthode non adversaire, stable et capable de préserver l'anatomie du patient sans données appariées.

2. Méthodologie : RAFM

Les auteurs proposent RAFM (Retrieval-Augmented Flow Matching), une approche basée sur le Flux Rectifié (Rectified Flow - RF) amélioré par une stratégie de récupération (retrieval).

A. Fondements du Flux Rectifié (RF)

Le RF modélise la traduction d'image comme un transport déterministe entre deux distributions (source CBCT et cible CT). Il apprend un champ de vitesse $v_\theta(x, t)$ qui guide l'interpolation linéaire entre un point source $x_0$ (CBCT) et un point cible $x_1$ (CT) via une équation différentielle ordinaire (ODE) :
$$\frac{dx_t}{dt} = v_\theta(x_t, t)$$
L'objectif est d'apprendre ce champ de vitesse en minimisant l'erreur entre la vitesse prédite et la vitesse cible ($x_1 - x_0$). Théoriquement, le RF ne nécessite pas d'appariement voxel-à-voxel, mais seulement que les marginales du couplage correspondent aux distributions source et cible.

B. Le Défi Pratique

Dans les petits ensembles de données médicaux et avec des tailles de lots (batch sizes) réduites, le couplage aléatoire ou local (au sein du lot) génère des paires d'extrémités sémantiquement incohérentes (ex: un CBCT de la tête couplé à un CT du bassin). Cela introduit du bruit dans le transport et dégrade la qualité de la traduction anatomique.

C. L'Innovation RAFM : Couplage Guidé par Récupération

Pour résoudre ce problème, RAFM remplace le couplage aléatoire par un couplage guidé par la récupération de caractéristiques (feature retrieval) :

1. Encodage de caractéristiques : Un encodeur DINOv3 (gelé) est utilisé pour extraire des embeddings de caractéristiques pour chaque tranche d'image.
2. Banque de Mémoire CT : Une banque de mémoire (memory bank) de capacité fixe stocke les paires (caractéristiques, image CT) des lots précédents (stratégie FIFO).
3. Construction de Paires Pseudo-Appariées : Pour chaque tranche CBCT du lot en cours, le système recherche dans la banque de mémoire la tranche CT la plus similaire en termes de similarité cosinus de leurs embeddings.
4. Entraînement : Ces paires "retrouvées" forment le couplage empirique $\rho_{retr}$ utilisé pour entraîner le modèle de flux.

Cette stratégie permet de construire des paires sémantiquement cohérentes (même région anatomique) sans utiliser d'identité de sujet ni d'appariement temporel réel, restant ainsi strictement dans le cadre "non apparié".

3. Contributions Clés

• Intégration du Flux Rectifié en Imagerie Médicale : Adaptation du RF pour la traduction CBCT-CT non appariée, offrant une alternative stable et non adversaire aux GANs et aux modèles de diffusion.
• Stratégie RAFM : Introduction d'un mécanisme de récupération global (via DINOv3 et une banque de mémoire) pour améliorer la qualité du couplage empirique dans des conditions de données limitées et de petits lots.
• Validation Rigoureuse : Évaluation sous un protocole "vraiment non apparié" (true-unpaired) au niveau du sujet, où aucune information de correspondance entre les sujets CBCT et CT n'est utilisée durant l'entraînement.

4. Résultats Expérimentaux

Les expériences ont été menées sur le benchmark SynthRAD2023 (bassin pelvien).

• Comparaison Quantitative : RAFM surpasse systématiquement les méthodes de l'état de l'art (CycleGAN, GcGAN, CUT, SynDiff, UNSB) sur toutes les métriques :
  • MAE (Erreur Absolue Moyenne) : 101.2 HU (meilleur résultat, contre 104.2 pour SynDiff).
  • FID (Fréchet Inception Distance) : 53.29 (indiquant une meilleure réalisme de la distribution, contre 62.91 pour UNSB).
  • SegScore (Consistance Anatomique) : 75.77% (supérieur aux autres méthodes), démontrant une excellente préservation des structures anatomiques.
• Analyse Qualitative : Les images générées par RAFM présentent une suppression plus propre des artéfacts et des structures anatomiques plus stables que les concurrents.
• Étude d'Ablation :
  • Le couplage aléatoire (RF standard) fonctionne mieux que la régression directe (U-Net classique) mais reste inférieur à RAFM.
  • L'augmentation de la taille de la banque de mémoire (jusqu'à $K=512$) améliore les performances, prouvant que la qualité du couplage est le facteur limitant principal.
  • RAFM se rapproche fortement des performances d'un modèle RF entraîné sur des données appariées (upper bound), notamment pour la préservation anatomique.

5. Signification et Impact

• Stabilité et Efficacité : RAFM offre une alternative robuste aux GANs, évitant l'instabilité de l'entraînement adversaire et les coûts de calcul élevés des modèles de diffusion (nécessitant de nombreux pas d'échantillonnage). L'inférence ne nécessite que 10 pas d'intégration ODE.
• Préservation Anatomique : En améliorant la qualité du couplage sans données appariées, RAFM résout le problème critique de la déformation anatomique souvent observé dans la traduction non appariée.
• Généralisabilité : La méthode est particulièrement adaptée aux scénarios cliniques réels où les données appariées sont rares ou inexistantes, ouvrant la voie à une utilisation plus fiable de la CBCT pour le calcul de dose en radiothérapie adaptative.

En résumé, RAFM démontre que l'amélioration de la qualité du couplage via la récupération de caractéristiques permet de tirer pleinement parti de la puissance théorique du Flux Rectifié pour des tâches de traduction d'images médicales complexes et non appariées.