Résumé technique : Entraînement par guidage supervisé pour les modèles de diffusion de dimension infinie

Énoncé du problème

Les problèmes inverses, en particulier ceux issus des équations aux dérivées partielles (EDP), nécessitent souvent d'inférer des paramètres inconnus existant dans des espaces fonctionnels de dimension infinie. Dans le cadre bayésien, l'objectif est d'échantillonner à partir d'une distribution a posteriori sur ces fonctions, conditionnée par des observations bruitées. Bien que les modèles de diffusion basés sur le score (SDM) aient récemment été étendus aux espaces fonctionnels de dimension infinie pour servir de priors expressifs, la théorie permettant de conditionner ces modèles afin d'échantillonner à partir de l'a posteriori reste un problème ouvert.

Les techniques existantes de conditionnement en dimension finie, telles que la transformation $h$ de Doob, reposent sur des décompositions de densité et des opérateurs de transition qui ne se traduisent pas directement en dimensions infinies. Les défis incluent l'inexistence de densités de Lebesgue, la singularité potentielle des mesures (théorème de Feldman-Hájek) lors du déplacement de mesures gaussiennes en dehors de leur espace de Cameron-Martin, et la difficulté d'assurer la différentiabilité des opérateurs de transition dans les espaces de Hilbert. De plus, la discrétisation immédiate des données pour appliquer des méthodes de dimension finie introduit des erreurs dépendantes de la résolution et limite la flexibilité du cadre.

Méthodologie

Fondement théorique : Transformation $h$ de dimension infinie

Les auteurs proposent un cadre rigoureux pour conditionner des modèles de diffusion de dimension infinie en utilisant une extension de la transformation $h$ de Doob. Ils établissent l'existence d'un processus de diffusion conditionnel dans deux configurations :

1. Configuration 1 : La distribution a priori est supportée à l'intérieur d'une boule de l'espace de Cameron-Martin du processus de Wiener utilisé dans l'équation différentielle stochastique (EDS) de diffusion.
2. Configuration 2 : Le prior est absolument continu par rapport à une mesure gaussienne, avec des bornes spécifiques sur la fonction de potentiel et son gradient.

Sous ces configurations, les auteurs prouvent que le processus conditionnel inversé dans le temps satisfait une équation différentielle stochastique (EDS) où la dérive est modifiée par un terme de guidage. Plus précisément, ils démontrent que la fonction de score conditionnelle, $s_y(t, x)$, se décompose en le score inconditionnel, $s(t, x)$, et un terme de guidage, $C\nabla \log h_y(t, x)$ :
$$s_y(t, x) = s(t, x) + C\nabla \log h_y(t, x)$$
Ici, $h_y$ est la transformation $h$ de dimension infinie définie via l'espérance conditionnelle du potentiel de vraisemblance.

Entraînement par guidage supervisé (SGT)

Puisque le terme de guidage $C\nabla \log h_y$ est intraitable en raison de sa dépendance à la loi de transition complète du processus de dimension infinie, les auteurs introduisent l'Entraînement par guidage supervisé (SGT). Il s'agit d'un objectif d'entraînement sans simulation conçu pour apprendre directement le terme de guidage.

• Objectif : La méthode minimise une perte de matching de score de débruitage. Au lieu de s'appuyer sur un contrôle stochastique optimal basé sur la simulation (qui souffre d'une variance élevée et nécessite l'échantillonnage de trajectoires), le SGT exploite le score inconditionnel connu $s(t, x)$ et la distribution conjointe des données et des observations $(X_0, Y)$.
• Fonction de perte : La perte est formulée pour minimiser la différence entre le terme de guidage appris $u(t, x, y)$ et le vrai terme de guidage. L'objectif est :
  $$\mathbb{E} \left[ \left\| \frac{X_t - e^{-t/2}X_0}{1 - e^{-t}} + s(t, X_t) + u(t, X_t, Y) \right\|^2_K \right]$$
  où $K$ est un espace de Hilbert choisi (par exemple, l'espace de Cameron-Martin).
• Paramétrisation : Le terme de guidage est paramétré à l'aide d'un réseau de neurones. Les auteurs explorent la préconditionnement de la sortie du réseau par l'opérateur de covariance $C$ à différents degrés ($k \in \{0, 1/2, 1\}$) pour imposer la régularité et s'aligner sur la structure de bruit de l'EDS.

Approximations et comparaisons

L'article analyse également les approches heuristiques existantes :

• Contrôle stochastique optimal (SOC) : Bien que théoriquement fondé, le SOC nécessite un entraînement basé sur la simulation, qui est coûteux en calcul et instable aux premiers stades de l'entraînement.
• Approximation de Tweedie (FunDPS) : Les auteurs étendent l'approximation FunDPS (qui utilise la formule de Tweedie pour estimer les données propres) à la Configuration 2 (continuité absolue), assouplissant les hypothèses restrictives des travaux antérieurs. Cependant, ils notent que cette approche engendre toujours une surcharge computationnelle due à la nécessité de rétropropager à travers le modèle de score à chaque étape de discrétisation.

Contributions clés

1. Extension théorique : L'article fournit la première preuve rigoureuse que les modèles de diffusion de dimension infinie peuvent être conditionnés via la transformation $h$ de Doob dans des configurations générales (support de Cameron-Martin ou continuité absolue), établissant la décomposition du score conditionnel en un score inconditionnel et un terme de guidage.
2. Entraînement par guidage supervisé (SGT) : Les auteurs proposent un algorithme d'entraînement novateur, sans simulation, pour apprendre le terme de guidage intraitable. Cette approche évite la variance élevée des méthodes basées sur la simulation et le coût computationnel de la différenciation à travers le modèle de score à chaque étape (comme requis par les approximations basées sur Tweedie).
3. Invariance de la résolution : En traitant le problème en dimensions infinies et en ne discrétisant qu'à l'étape finale, le cadre offre une approche invariante par la résolution pour le conditionnement, évitant la dégradation des performances souvent observée lors de l'application d'architectures de dimension finie à des résolutions plus élevées.
4. Généralisation de FunDPS : Le travail démontre que les approximations basées sur Tweedie (FunDPS) sont valables dans le cadre moins restrictif de la Configuration 2 (continuité absolue), élargissant ainsi leur applicabilité.

Résultats expérimentaux

Les auteurs évaluent le SGT sur trois problèmes inverses bayésiens dans des espaces fonctionnels :

1. Observations clairsemées : Reconstruction d'une fonction 1D à partir d'observations ponctuelles bruitées.
2. Équation de la chaleur : Reconstruction de la distribution initiale de température d'une équation de la chaleur à partir d'observations bruitées à un temps ultérieur.
3. Inpainting de formes : Reconstruction de formes (chiffre "3" issu de MNIST) à partir de coefficients de Fourier partiels (descripteurs de Fourier elliptiques).

Constats :

• Performance : Le SGT surpasse systématiquement la ligne de base heuristique FunDPS en termes d'erreur quadratique moyenne (RMSE) et de score d'énergie (ES).
• Comparaison avec la diffusion conditionnelle : Le SGT atteint une précision comparable à celle d'un modèle de diffusion "conditionnel" entraîné à partir de zéro (qui sert de borne supérieure), malgré le fait que le modèle conditionnel soit significativement plus grand et entraîné pendant plus d'étapes.
• Robustesse : Dans des expériences avec des modèles de base inconditionnels imparfaits (initialisés aléatoirement ou partiellement entraînés), le SGT démontre une robustesse supérieure, récupérant des échantillons réalistes là où FunDPS échoue complètement.
• Efficacité : Le SGT nécessite moins de paramètres et d'étapes d'entraînement que l'entraînement d'un modèle conditionnel complet à partir de zéro, car il n'entraîne que le réseau de guidage tout en réutilisant un modèle inconditionnel pré-entraîné.

Importance et affirmations

L'article prétend offrir la première méthode dans l'espace des fonctions pour le fine-tuning de modèles de diffusion entraînés afin d'échantillonner précisément à partir d'une distribution a posteriori. L'importance réside dans le comblement du fossé entre les exigences théoriques de l'inférence bayésienne de dimension infinie et la mise en œuvre pratique des modèles de diffusion.

Les auteurs soulignent que leur cadre permet :

• Invariance de la résolution : La capacité de générer des échantillons à des résolutions différentes de celles des données d'entraînement sans dégradation des performances.
• Stabilité : Un régime d'entraînement stable même lorsque le modèle inconditionnel sous-jacent est imparfait.
• Flexibilité : Des options alternatives pour la discrétisation, l'entraînement et les fonctions de perte qui ne sont pas contraintes par des hypothèses de densité de dimension finie.

Le travail conclut que l'apprentissage explicite du terme de guidage de dimension infinie produit une précision supérieure par rapport aux approximations heuristiques, offrant une voie rigoureuse pour appliquer les modèles de diffusion à des problèmes inverses complexes en science et en ingénierie.