Flow-Based Generative Modeling for Optimizing Sampling Policies in Compressed Sensing Applications

Résumé Technique : Modélisation Générative Basée sur le Flux pour l'Optimisation des Politiques d'Échantillonnage en Récupération Compressée

Énoncé du Problème
Les applications modernes en traitement du signal et en imagerie médicale nécessitent souvent l'acquisition de signaux de grande dimension sous des contraintes de ressources strictes. La théorie classique de l'échantillonnage dicte qu'une reconstruction précise nécessite des mesures proportionnelles à la dimension ambiante du signal, une exigence qui est fréquemment impraticable. La Récupération Compressée (Compressed Sensing - CS) répond à cela en permettant la récupération de signaux parcimonieux à partir de moins de mesures, à condition que l'opérateur de mesure satisfasse des conditions spécifiques (par exemple, l'incohérence). Cependant, la qualité de la reconstruction dépend fortement du choix du schéma d'échantillonnage. Trouver un masque de sous-échantillonnage optimal est un problème combinatoire NP-difficile car le masque d'échantillonnage est discret et non différentiable, ce qui empêche une optimisation directe de bout en bout. De plus, l'espace de recherche croît exponentiellement avec le nombre d'emplacements de capteurs, rendant l'exploration exhaustive infaisable. Les approches génératives existantes, telles que le Deep Probabilistic Subsampling (DPS) ou l'Active Diffusion Subsampling (ADS), produisent souvent des masques fixes ou souffrent d'une latence de calcul élevée, limitant leur adaptabilité et leur utilité clinique en temps réel.

Méthodologie
Les auteurs proposent un cadre génératif basé sur le flux, sensible aux tâches, qui reformule le paradigme d'entraînement conventionnel du Flow Matching (FM). Au lieu de modéliser directement la variété des données, le modèle opère sur la variété de l'espace d'échantillonnage pour apprendre des masques de sous-échantillonnage optimaux.

• Flow Matching sur la Variété d'Échantillonnage : Le cadre utilise une équation différentielle ordinaire (ODE) neuronale pour transporter des points d'une distribution de bruit gaussien standard ($p_0$) vers une distribution cible d'opérateurs d'échantillonnage ($p_1$). Le modèle apprend un champ de vecteurs dépendant du temps $v_\phi(x, t)$ qui définit une trajectoire générative.
• Optimisation Conditionnée par la Tâche : Contrairement au FM standard qui repose sur un couplage linéaire entre le bruit et les données, ce cadre couple le flux génératif avec un solveur de tâche CS en aval ($f_\theta$). L'objectif d'entraînement minimise une fonction de perte composite combinant la perte de régression FM standard ($L_{FM}$) et une perte spécifique à la tâche ($L_{task}$), pondérée par des paramètres d'incertitude apprenables ($\sigma_1, \sigma_2$). Cela permet au modèle d'optimiser conjointement la politique d'échantillonnage et la performance de la tâche.
• Porte Sigmoïde Différentiable : Pour traiter la non-différentiabilité des masques binaires discrets, la sortie continue du modèle de flux ($x_t$) est passée à travers une porte sigmoïde personnalisée et différentiable. Cette porte emploie une fonction de raideur dépendante du temps $s(t)$ et un seuil de percentile $\tau$ pour produire des masques souples ($g(x_t) \in (0, 1)$). À mesure que le pas de temps $t$ augmente, la transition devient plus abrupte, régularisant efficacement l'entropie du processus de génération du masque.
• Modes d'Inférence : Le cadre prend en charge deux modes opérationnels :
  1. Du Aléatoire à l'Optimisé : Des mesures aléatoires initiales sont effectuées, les caractéristiques sont extraites, puis le modèle de flux génère un masque optimisé pour les mesures suivantes. Cela convient aux scénarios où les mesures sont peu coûteuses ou les signaux sont statiques (ex. : MNIST, CelebA).
  2. Échantillonnage Adaptatif : Des données de l'espace k à basse fréquence (LF) sont acquises en premier. Un extracteur de caractéristiques conditionne le modèle de flux pour déterminer le schéma d'échantillonnage optimal pour les bandes de haute fréquence restantes. Ce mode est conçu pour les mesures coûteuses, spécifiquement l'accélération IRM.

Contributions Clés

• Nouveau Paradigme : L'article introduit une preuve de concept où le Flow Matching est entraîné au sein de la variété des opérateurs d'échantillonnage plutôt que dans l'espace du signal, spécifiquement adapté aux tâches de CS en aval.
• Politique Continue et Adaptative : Le cadre fournit une paramétrisation continue et temporellement continue du schéma d'échantillonnage, offrant une évolution fluide des structures grossières vers des raffinements spécifiques à la tâche, contrairement aux méthodes de génération de masques statiques.
• Optimisation Conjointe : La méthode permet une optimisation conjointe de bout en bout de la stratégie d'échantillonnage et de la tâche de reconstruction/classification, en exploitant l'a priori de la tâche pour guider le processus génératif.
• Cadre Unifié : Il offre une approche flexible applicable à divers problèmes inverses, incluant la classification d'images, la reconstruction d'images et l'accélération de l'IRM, sans nécessiter de réentraînement pour différents opérateurs de mesure.

Résultats Expérimentaux
Le cadre a été évalué sur MNIST, CelebA et le jeu de données fastMRI knee :

• Classification MNIST : L'approche basée sur le flux a amélioré la précision de la classification de plus de 20 % par rapport aux classificateurs entraînés sur des masques aléatoires. Bien qu'elle soit légèrement en retrait par rapport au modèle A-DPS pour certains ratios d'échantillonnage fixes, elle a démontré une plus grande polyvalence sur une large gamme de ratios (1 % à 8 %) sans nécessiter d'entraînement spécifique au ratio.
• Reconstruction d'Image (CelebA & MNIST) : La méthode a atteint des performances de pointe en reconstruction d'image. Sur le jeu de données CelebA, elle a atteint un rapport signal sur bruit de crête (PSNR) de 25,17 dB à un taux de sous-échantillonnage de 5 %. Elle a surpassé la méthode de pointe ADS, particulièrement lors de taux d'échantillonnage agressifs (faibles).
• Accélération IRM : Dans les expériences fastMRI, le cadre a obtenu un indice de similitude structurelle (SSIM) de 0,7163 à une accélération de 8×, surpassant toutes les méthodes de référence (LOUPE, DPS, ADS, A-DPS) à ce facteur d'accélération élevé. À une accélération de 4×, ses performances étaient comparables aux autres méthodes.
• Efficacité de Calcul : Le temps d'inférence pour générer un masque d'échantillonnage était d'environ 127–130 ms sur une NVIDIA GeForce RTX 4070. C'est nettement plus rapide que le temps d'acquisition par ligne de l'espace k en IRM (500–2000 ms) et drastiquement plus rapide que le temps d'acquisition par ligne rapporté pour la méthode ADS (3040 ms), suggérant une faisabilité pour un déploiement clinique en temps réel.

Signification et Revendications
Les auteurs affirment que ce travail établit la faisabilité fondamentale de l'utilisation de modèles génératifs basés sur le flux conditionnés par des tâches pour optimiser les politiques d'échantillonnage en récupération compressée. L'étude souligne que :

1. Le Conditionnement par la Tâche est Efficace : L'intégration d'un a priori de tâche en aval directement dans le processus d'entraînement du Flow Matching améliore significativement les performances dans les problèmes inverses.
2. Contraintes de la Variété d'Échantillonnage : L'espace d'échantillonnage possède des contraintes structurelles intrinsèques (binarité, parcimonie) qui peuvent être exploitées pour améliorer l'efficacité de l'entraînement par rapport à la modélisation de la variété d'image naturelle.
3. Viabilité Clinique : La faible surcharge de calcul et la nature adaptative du cadre en font un candidat prometteur pour les applications d'imagerie médicale réelles, particulièrement pour réduire les temps d'examen IRM et améliorer la qualité des données diagnostiques.
4. Généralisabilité : L'approche offre une stratégie unifiée et flexible pour concevoir des schémas de détection pilotés par les données et les tâches, applicables à une large gamme de problèmes inverses au-delà des tâches spécifiques testées.

L'article se positionne comme une étude de preuve de concept, reconnaissant que bien qu'il atteigne des résultats de pointe dans des tâches de reconstruction spécifiques, la contribution principale est la démonstration d'un nouveau cadre flexible pour apprendre des politiques d'échantillonnage qui peuvent être adaptées à divers problèmes inverses.