Flow Matching with Optimized Subclass Priors for Medical Image Augmentation

Résumé technique : Flow Matching avec des Priors de Sous-classes Optimisés pour l'Augmentation d'Images Médicales

Énoncé du Problème
Les ensembles de données d'imagerie médicale se caractérisent par de graves déséquilibres à longue traîne, où les affections courantes dominent tandis que les pathologies rares et cliniquement critiques (par exemple, les petits pneumothorax, les bronchectasies) sont sous-représentées. Ce déséquilibre amène les classificateurs discriminatifs à surajuster aux classes majoritaires et à échouer sur les affections rares. Bien que l'augmentation de données générative offre une solution potentielle, le flow matching (FM) conditionnel standard et les modèles de diffusion rencontrent deux difficultés cumulatives lorsqu'ils sont entraînés sur des étiquettes de maladies grossières :

1. Agrégation Multi-modale : Une étiquette diagnostique unique agrège souvent des sous-types divers, des protocoles d'acquisition et des comorbidités. Cela crée des conditionnelles multi-modales qui biaisent les générateurs vers les sous-modes dominants, tronquant efficacement les variations rares mais cliniquement valides.
2. Géométrie de Transport Inefficace : Le fait de tirer toutes les trajectoires de génération à partir d'une source gaussienne standard partagée force des sous-populations spécifiques à emprunter des chemins de transport disproportionnellement longs. Cela gonfle la difficulté de régression et dégrade la fidélité de la génération des classes de la queue de distribution.

Méthodologie
Les auteurs proposent une stratégie hors ligne qui introduit des priors informatifs à deux niveaux pour affiner le processus génératif sans modifier la boucle d'entraînement principale du modèle de flow matching.

1. Induction de Sous-classes via Modélisation par Mélange Gaussien (GMM) :

   • La méthode partitionne chaque étiquette grossière $c$ en sous-modes cohérents (sous-classes $k$) dans l'espace latent du modèle génératif.
   • En utilisant un VAE figé (FLUX1.dev), des latents déterministes sont calculés. Pour chaque classe $c$, un modèle de mélange gaussien diagonal est ajusté aux latents résiduels (après soustraction du centre de classe $\mu_c$).
   • Le nombre de composantes $K_c$ est sélectionné en utilisant le Critère d'Information Bayésien Étendu (EBIC). Les échantillons reçoivent des affectations rigides aux sous-classes, créant une partition avec une variance conditionnelle réduite.
   • Théoriquement, les auteurs prouvent que le conditionnement sur ces étiquettes affinées ne peut pas augmenter le risque optimal de flow matching de Bayes ; l'amélioration est égale à la variabilité attendue entre les sous-classes du déplacement moyen conditionnel.

2. Optimisation de la Source Dépendante de la Sous-classe :

   • Au lieu d'une source gaussienne partagée, la méthode apprend une distribution source spécifique $\pi_0(\cdot | c, k)$ pour chaque sous-classe. Ces sources sont paramétrées comme des gaussiennes apprenables $N(\mu_{c,k}, \text{diag}(\sigma^2_{c,k}))$ qui recentrent et re-échelonnent la distribution de départ.
   • Cela raccourcit les trajectoires de transport typiques et réduit la dispersion au sein de la sous-classe.
   • Contrôle Géométrique : Pour éviter les solutions dégénérées (telles que l'inflation de norme ou des chemins trop concentrés), les auteurs optimisent un objectif proxy hors ligne qui impose :
     • Compacité Directionnelle : Encourager les directions de déplacement normalisées à s'aligner modérément avec des prototypes apprenables.
     • Plafonnement de la Longueur de Trajet : Imposer une pénalité fixe aux valeurs aberrantes en utilisant un plafond spécifique à la sous-classe ($r_{cap}$), défini comme le 99e percentile des distances au centre de classe. Cela empêche le modèle de contourner la perte en gonflant les normes de déplacement.
     • Régularisation d'Échelle : Régulariser les échelles logarithmiques des distributions sources vers une référence unitaire.

Contributions Clés

1. Formulation Conditionnelle par Sous-classe : Une approche novatrice pour la génération d'images médicales à longue traîne qui partitionne les conditionnelles multi-modales via des affectations rigides induites par GMM, réduisant efficacement l'ambiguïté conditionnelle.
2. Sources Dépendantes de la Sous-classe Apprenables : L'introduction de priors informatifs qui remodelent le transport conditionnel en recentrant et re-échelonnant la distribution source par sous-classe, réduisant ainsi la dispersion au sein de la sous-classe.
3. Objectif Régularisé Géométriquement : Un objectif proxy pratique qui équilibre une compacité angulaire modérée avec un contrôle explicite des valeurs aberrantes de longueur de trajet, empêchant les raccourcis d'inflation de norme et l'effondrement directionnel.
4. Validation Empirique : Démonstration d'une meilleure couverture de la queue de distribution et de performances de classification à longue traîne en aval sur des benchmarks de radiographies thoraciques (MIMIC-LT, NIH-LT) et de scanners CT (CT-RATE).

Résultats Expérimentaux
La méthode a été évaluée sur trois benchmarks d'imagerie médicale à longue traîne :

• Classification en Aval : La méthode proposée a constamment obtenu les scores de précision équilibrée (bAcc) et de macro-F1 les plus élevés sur les trois ensembles de données par rapport aux bases de référence sans augmentation et à d'autres stratégies d'augmentation (FM Vanilla, CBFM, CPD). Notamment, c'était la seule stratégie à améliorer de manière fiable les performances à la fois sur les modalités Rayons X et CT.
• Qualité Générative : Sur MIMIC-LT, la méthode a obtenu la meilleure distance Fréchet Inception (FID) et le meilleur score de rapport Inception (IRS), indiquant des gains simultanés en fidélité et en diversité. Sur CT-RATE, bien que le FID ait été légèrement plus élevé que celui d'une base de référence, la méthode a maintenu un IRS élevé et une classification en aval supérieure, suggérant une meilleure couverture distributionnelle plutôt qu'une simple mémorisation.
• Études d'Ablation : Les résultats ont confirmé que les deux étapes du pipeline (sous-classification et optimisation de source) sont nécessaires. La sous-classification seule a amélioré les performances, mais la combinaison a produit les meilleurs résultats. L'optimisation sans sous-classification a parfois dégradé les performances en moyennant des modes conflictuels.
• Validité des Sous-classes : L'analyse a montré que les sous-classes induites capturent une structure pertinente pour la pathologie plutôt que simplement des métadonnées d'acquisition. Des sondes de confondants ont indiqué que les métadonnées étaient détectables mais non dominantes, et la pureté kNN des sous-classes apprises dépassait significativement celle des partitions aléatoires.

Signification et Revendications
L'article revendique que la conception de la source et l'affinement conditionnel sont des leviers efficaces pour une augmentation générative robuste dans des contextes cliniques à longue traîne. En partitionnant les étiquettes hétérogènes et en apprenant des sources dépendantes de la sous-classe, la méthode réduit l'ambiguïté conditionnelle et stabilise les trajectoires de transport à un coût computationnel négligeable (l'ensemble du processus est hors ligne et prend moins d'une heure de GPU par jeu de données). Les auteurs positionnent cela comme une stratégie fondée sur des principes qui améliore la fidélité et la diversité de la génération des classes de la queue, conduisant à des classificateurs en aval plus fiables pour les affections médicales rares. Le travail souligne que la génération conditionnelle naïve échoue souvent sur les sous-populations rares en raison de problèmes de géométrie de transport, qui peuvent être atténués par des priors structurés.