An Unconditional Representation of the Conditional Score in Infinite-Dimensional Linear Inverse Problems

Résumé technique : Une représentation inconditionnelle du score conditionnel dans les problèmes inverses linéaires de dimension infinie

Énoncé du problème
Les modèles de diffusion basés sur le score (SDM) sont devenus un outil puissant pour l'échantillonnage à partir de distributions a posteriori dans les problèmes inverses bayésiens, offrant des a priori flexibles qui surpassent les régularisateurs traditionnels conçus à la main. Cependant, les méthodes existantes pour l'échantillonnage a posteriori font face à un compromis critique entre l'efficacité computationnelle et l'évolutivité.

1. Les méthodes basées sur la correction (par exemple, l'ajout de gradients de vraisemblance au score ou l'imposition de la cohérence des données) nécessitent des évaluations répétées de l'opérateur direct $A$ au cours du processus d'échantillonnage. Cela crée un goulot d'étranglement computationnel prohibitif pour les problèmes à grande échelle et de haute dimension.
2. Les méthodes d'entraînement conditionnel direct apprennent une fonction de score conditionnée aux mesures ($s(x, t; \mu_y)$). Bien que cela évite les évaluations directes lors de l'échantillonnage, cela augmente considérablement la dimensionnalité d'entrée du réseau de neurones (de $x$ à $(x, y)$), nécessitant des données d'entraînement extensives et des architectures complexes pour approximer avec précision le score conditionnel, en particulier dans des contextes de dimension infinie.

L'article aborde le défi de réaliser un échantillonnage a posteriori efficace et évolutif dans les problèmes inverses linéaires de dimension infinie sans encourir le coût en ligne élevé des évaluations du modèle direct ni la complexité d'échantillonnage élevée de l'entraînement conditionnel direct.

Méthodologie : Représentation inconditionnelle du score conditionnel (UCoS)
Les auteurs proposent UCoS, un cadre qui déplace la charge computationnelle de la conditionnement de la phase d'échantillonnage en ligne vers une phase d'entraînement hors ligne. L'idée centrale est que pour les problèmes inverses linéaires avec un bruit gaussien, la fonction de score conditionnelle peut être dérivée exactement à partir d'une fonction de score inconditionnelle d'un processus de diffusion modifié et dépendant de la tâche, via des transformations affines.

• Fondement théorique :

  • La méthode opère dans un espace de Hilbert séparable $\mathcal{H}$, assurant l'invariance de la discrétisation.
  • Au lieu d'entraîner un réseau pour approximer $s(x, t; \mu_y)$, les auteurs définissent un nouveau processus de diffusion $\tilde{X}_t$ dont la distribution est un mélange de l'a priori et d'un terme gaussien dépendant de l'opérateur direct $A$.
  • Ils démontrent (Théorème 3.7) que le vrai score conditionnel $s(x, t; \mu_y)$ est lié au score inconditionnel $\tilde{s}$ de ce nouveau processus par l'identité :
    $$s(x, t; \mu_y) = \lambda(t) \left( r(\xi_t(x, y)) - e^{-t/2}x \right)$$
    où $r$ est une version transformée du score inconditionnel, $\lambda(t)$ est un scalaire dépendant du temps, et $\xi_t(x, y)$ est une transformation affine impliquant la mesure $y$ et l'opérateur adjoint $A^*$.
  • Crucialement, le terme $A^*\Gamma^{-1}y$ (où $\Gamma$ est la covariance du bruit) peut être calculé une fois avant le début de l'échantillonnage. Le processus d'échantillonnage subséquent ne nécessite que l'évaluation du réseau entraîné $r_\theta$ et de simples additions vectorielles, éliminant le besoin d'appliquer davantage $A$ ou $A^*$ lors des étapes de diffusion inverse.

• Implémentation :

  • Phase hors ligne : Un réseau de neurones (spécifiquement un Opérateur de Réseau de Fourier, FNO) est entraîné pour approximer la fonction $r$. L'entraînement utilise un objectif d'appariement de score de débruitage sur le processus modifié $\tilde{X}_t$. Cette phase nécessite l'accès à $A$ et $A^*$ mais aucune donnée de mesure spécifique $y$.
  • Phase en ligne : Des échantillons a posteriori sont générés en résolvant l'EDS à temps inverse en utilisant le score transformé. Le processus est « sans matrice », ne nécessitant que le décalage précalculé $A^*\Gamma^{-1}y$ et l'évaluation du réseau.

Contributions clés

1. Identité exacte en dimension infinie : L'article établit une identité théorique rigoureuse (Théorème 3.7) reliant le score conditionnel à un score inconditionnel dépendant de la tâche via des transformations affines dans des espaces de Hilbert de dimension infinie. Cela généralise le principe au-delà de discrétisations spécifiques de dimension finie.
2. Analyse de convergence : Les auteurs fournissent une borne d'erreur quantitative (Théorème 4.1) pour le processus génératif, caractérisant explicitement les contributions d'erreur provenant de la perte d'entraînement, de la discrétisation numérique et de l'initialisation. Cela justifie la convergence de la méthode vers la mesure a posteriori vraie.
3. Évolutivité et efficacité : En découplant le conditionnement de la boucle d'échantillonnage, UCoS atteint l'efficacité computationnelle des méthodes inconditionnelles (aucune évaluation directe pendant l'échantillonnage) tout en conservant la précision des méthodes conditionnelles. Il évite le fléau de la dimensionnalité associé à l'entraînement de réseaux sur l'espace conjoint $(x, y)$.

Résultats expérimentaux
Les auteurs valident UCoS sur trois problèmes inverses : l'inpainting d'images, la tomographie assistée par ordinateur (CT) à vues éparses et le défloutage d'images.

• Inpainting (A priori gaussien) : Dans un contexte où l'a posteriori est gaussien et les scores peuvent être calculés sous forme fermée, UCoS se comporte de manière comparable à la méthode conditionnelle exacte et à d'autres références, confirmant l'identité théorique.
• Imagerie CT et défloutage (A priori appris) : En utilisant respectivement les ensembles de données LIDC-IDRI et CelebA, UCoS est comparé à SDE ALD, à l'échantillonnage a posteriori par diffusion (DPS), aux méthodes basées sur la projection (Proj) et à l'entraînement conditionnel direct.
  • Précision : UCoS produit des échantillons avec un biais et un écart type comparables ou meilleurs que les méthodes inconditionnelles et nettement meilleurs que la méthode conditionnelle lors de l'utilisation d'architectures de faible complexité (par exemple, 32 nœuds par couche).
  • Efficacité : UCoS correspond à la vitesse d'échantillonnage de la méthode conditionnelle (toutes deux évitent les évaluations directes) mais surpasse nettement les méthodes basées sur la correction (DPS, Proj, SDE ALD) qui nécessitent des passes directes répétées.
  • Robustesse : Les performances de la méthode conditionnelle se dégradent notablement avec moins de paramètres de réseau en raison de la complexité accrue de l'approximation de la distribution conjointe, tandis que UCoS maintient une haute qualité avec des architectures plus simples.

Signification et affirmations
L'article affirme que UCoS offre une solution « le meilleur des deux mondes » : il élimine le goulot d'étranglement computationnel des évaluations de modèles directs lors de l'échantillonnage (une limitation des méthodes basées sur la correction) tout en évitant la complexité d'échantillonnage élevée et les exigences architecturales de l'entraînement conditionnel direct.

Les auteurs soulignent que leur approche est invariante par discrétisation, la rendant adaptée aux problèmes inverses de dimension infinie où les garanties de performance échouent souvent dans les approximations de dimension finie de haute dimension. Ils reconnaissent que la méthode est dépendante de la tâche, nécessitant un ré-entraînement si l'opérateur direct $A$ change, et que la correspondance exacte repose sur la linéarité du modèle direct et la gaussianité de la vraisemblance (permettant l'étape de « complétion du carré »). Le travail est présenté comme une preuve de concept pour une nouvelle classe d'échantillonneurs basés sur la diffusion qui déplacent la complexité du conditionnement vers la phase d'entraînement sans erreur d'approximation dans la formulation théorique.