Transport Optimal et Convergence des Modèles de Diffusion

Cet article étudie les propriétés de convergence des modèles génératifs basés sur le score (SGMs) à travers le prisme du Transport Optimal. Nous modélisons le processus de génération comme une Équation Différentielle Stochastique (EDS) inverse temporellement. En utilisant le théorème de Girsanov et l'inégalité de Talagrand, nous établissons une borne supérieure explicite sur la distance de Wasserstein-2 entre la distribution de données réelle et la distribution générée, en fonction de l'erreur d'approximation du score ($L^2$-score matching error).

1. Introduction

Les modèles de diffusion probabilistes ont émergé comme l'état de l'art en modélisation générative. Le principe repose sur deux processus : un processus direct (Forward) qui détruit l'information en ajoutant du bruit, et un processus inverse (Reverse) qui reconstruit les données à partir du bruit.

Soit $p_{\text{data}}$ la distribution des données sur $\mathbb{R}^d$. Le défi mathématique fondamental est de quantifier la qualité de l'échantillonnage. Si $p_{\theta}$ est la distribution générée par le modèle paramétré $\theta$, nous cherchons à borner la distance de Wasserstein d'ordre 2, notée $W_2(p_{\text{data}}, p_{\theta})$.

Contrairement aux approches basées sur la divergence de Kullback-Leibler (KL) ou la distance Variationnelle Totale (TV), la métrique $W_2$ prend en compte la géométrie sous-jacente de l'espace des données, ce qui est crucial pour les variétés de basse dimension dans des espaces de haute dimension (hypothèse de la variété).

2. Cadre Théorique : Dynamique de Langevin et EDS

2.1 Le Processus Direct (Forward)

Nous considérons un processus de diffusion continu $\{ \mathbf{x}_t \}_{t \in [0, T]}$ régi par l'EDS d'Ornstein-Uhlenbeck :

$$ d\mathbf{x}_t = -\frac{1}{2}\mathbf{x}_t dt + d\mathbf{w}_t, \quad \mathbf{x}_0 \sim p_{\text{data}} $$

où $\mathbf{w}_t$ est un mouvement brownien standard dans $\mathbb{R}^d$. La densité marginale de $\mathbf{x}_t$ est notée $p_t$. Pour $T \to \infty$, $p_T$ converge vers la distribution stationnaire $\pi(\mathbf{x}) = \mathcal{N}(\mathbf{0}, \mathbf{I}_d)$.

2.2 Le Processus Inverse (Reverse)

D'après le théorème d'Anderson (1982), le processus de retournement temporel (qui va de $t=T$ à $t=0$) est également un processus de diffusion, régi par :

$$ d\mathbf{x}_t = \left[ -\frac{1}{2}\mathbf{x}_t - \nabla_{\mathbf{x}} \log p_t(\mathbf{x}_t) \right] dt + d\bar{\mathbf{w}}_t $$

où $dt$ représente ici un incrément de temps négatif et $\bar{\mathbf{w}}_t$ est un mouvement brownien standard pour le temps inversé. Le terme $\nabla_{\mathbf{x}} \log p_t(\mathbf{x})$ est la fonction de score (Score function).

2.3 Le Modèle Génératif Approché

En pratique, nous n'avons pas accès à $\nabla \log p_t$. Nous entraînons un réseau de neurones $s_\theta(\mathbf{x}, t)$ pour l'approximer via la minimisation de la Score Matching Loss :

$$ \mathcal{L}(\theta) = \mathbb{E}_{t \sim \mathcal{U}[0,T], \mathbf{x}_t \sim p_t} \left[ \| s_\theta(\mathbf{x}_t, t) - \nabla \log p_t(\mathbf{x}_t) \|^2_2 \right] $$

Le processus génératif réel, noté $\mathbf{y}_t$, suit alors l'EDS approchée partant de $\mathbf{y}_T \sim \pi$ :

$$ d\mathbf{y}_t = \left[ -\frac{1}{2}\mathbf{y}_t - s_\theta(\mathbf{y}_t, t) \right] dt + d\bar{\mathbf{w}}_t $$

Soit $p_\theta$ la loi de $\mathbf{y}_0$ (la donnée générée finale).

3. Résultats Principaux

Nous cherchons à lier l'erreur d'apprentissage du score, $\epsilon_{\text{score}}$, à la distance $W_2(p_{\text{data}}, p_\theta)$.

Définition 1 (Distance de Wasserstein $W_2$) Pour deux mesures de probabilité $\mu, \nu$ sur $\mathbb{R}^d$ ayant des moments d'ordre 2 finis :

$$ W_2^2(\mu, \nu) = \inf_{\gamma \in \Pi(\mu, \nu)} \int_{\mathbb{R}^d \times \mathbb{R}^d} \|x - y\|^2 d\gamma(x, y) $$

où $\Pi(\mu, \nu)$ est l'ensemble des couplages de $\mu$ et $\nu$.

Théorème 1 (Borne Supérieure de Convergence) Soit $p_{\text{data}}$ une distribution satisfaisant l'inégalité de Talagrand $T_2(C)$ avec une constante $C > 0$. Supposons que l'approximation du score soit bornée uniformément en temps par :

$$ \sup_{t \in [0, T]} \mathbb{E}_{\mathbf{x} \sim p_t} [ \| s_\theta(\mathbf{x}, t) - \nabla \log p_t(\mathbf{x}) \|^2 ] \le \varepsilon^2 $$

Alors, la distance de Wasserstein entre les données réelles et les données générées satisfait :

$$ W_2(p_{\text{data}}, p_\theta) \le \sqrt{2C \cdot \mathcal{D}_{KL}(p_{\text{data}} || p_\theta)} \le \sqrt{C T} \cdot \varepsilon + \mathcal{O}(e^{-T}) $$

Le terme $\mathcal{O}(e^{-T})$ représente l'erreur due à l'initialisation imparfaite (la convergence vers la gaussienne n'est jamais totale à temps fini), souvent négligeable pour $T$ grand.

4. Éléments de Preuve

La preuve repose sur l'utilisation du théorème de Girsanov pour estimer la divergence KL entre les mesures de chemin (Path Measures), puis l'utilisation de l'inégalité de transport.

4.1 Changement de Mesure et Théorème de Girsanov

Notons $\mathbb{P}$ la mesure de chemin induite par le processus inverse exact, et $\mathbb{Q}$ la mesure induite par le processus génératif (avec le score approché). La dérivée de Radon-Nikodym entre les trajectoires sur $[0, T]$ est donnée par l'exponentielle de Girsanov :

$$ \frac{d\mathbb{P}}{d\mathbb{Q}} = \exp\left( \int_0^T \langle \nabla \log p_t(\mathbf{x}_t) - s_\theta(\mathbf{x}_t, t), d\bar{\mathbf{w}}_t \rangle - \frac{1}{2} \int_0^T \| \nabla \log p_t(\mathbf{x}_t) - s_\theta(\mathbf{x}_t, t) \|^2 dt \right) $$

4.2 Divergence KL sur les Chemins

La divergence de Kullback-Leibler entre les distributions de chemins complets est :

$$ \mathcal{D}_{KL}(\mathbb{P}_{[0,T]} || \mathbb{Q}_{[0,T]}) = \mathbb{E}_{\mathbb{P}} \left[ \log \frac{d\mathbb{P}}{d\mathbb{Q}} \right] $$

En prenant l'espérance, le terme stochastique (intégrale d'Itô) s'annule (sous conditions de martingale). Il reste le terme de dérive quadratique. De plus, par l'inégalité de traitement des données (Data Processing Inequality), la divergence sur les marginaux à $t=0$ est inférieure ou égale à la divergence sur les chemins complets :

$$ \mathcal{D}_{KL}(p_{\text{data}} || p_\theta) \le \mathcal{D}_{KL}(\mathbb{P}_{[0,T]} || \mathbb{Q}_{[0,T]}) = \frac{1}{2} \int_0^T \mathbb{E}_{p_t} \left[ \| \nabla \log p_t(\mathbf{x}_t) - s_\theta(\mathbf{x}_t, t) \|^2 \right] dt $$

Nous reconnaissons ici l'intégrale temporelle de l'erreur de score matching (la divergence de Fisher pondérée). Si cette erreur est bornée par $\varepsilon^2$ uniformément, alors :

$$ \mathcal{D}_{KL}(p_{\text{data}} || p_\theta) \le \frac{1}{2} T \varepsilon^2 $$

4.3 Passage à Wasserstein (Inégalité de Talagrand)

Pour conclure, nous devons lier la divergence KL à la distance $W_2$. Une mesure $\mu$ satisfait l'inégalité de Talagrand $T_2(C)$ si pour toute mesure $\nu$ :

$$ W_2^2(\nu, \mu) \le 2C \cdot \mathcal{D}_{KL}(\nu || \mu) $$

En supposant que $p_{\text{data}}$ vérifie cette propriété (ce qui est vrai si $p_{\text{data}}$ est, par exemple, fortement log-concave, ou une perturbation bornée d'une gaussienne), nous combinons les résultats :

$$ W_2(p_{\text{data}}, p_\theta) \le \sqrt{2C \cdot \frac{1}{2} T \varepsilon^2} = \varepsilon \sqrt{CT} $$

(Note : Une analyse plus fine traiterait l'erreur d'initialisation à $t=T$ via l'inégalité triangulaire de Wasserstein, ajoutant le terme exponentiellement décroissant mentionné dans le Théorème).

Nous avons démontré formellement que la minimisation de l'erreur de score matching (Fisher Divergence) entraîne une minimisation de la distance de Wasserstein $W_2$ pour la génération. Ce résultat justifie mathématiquement l'efficacité des modèles de diffusion : apprendre le champ de vecteurs local (le score) garantit le transport optimal global de la mesure de probabilité.

Les travaux futurs pourraient explorer le relâchement de l'hypothèse de Talagrand, notamment pour les distributions à support compact ou les variétés de dimension inférieure ("Manifold Hypothesis"), où la constante $C$ peut devenir problématique.