Résumé Technique : URGE – Approximation Simple et Inférence sans Dérivée à l'Exécution pour les Modèles de Diffusion via Monte Carlo Séquentiel sur les Mesures de Trajectoire

1. Énoncé du Problème

Les modèles de diffusion modernes et les méthodes génératives basées sur le flot sont généralement entraînés pour capturer la distribution des données $p_{\text{data}}$. Cependant, le déploiement nécessite souvent d'orienter la génération vers des objectifs spécifiques en aval (par exemple, validité physique, alignement texte-image) sans réentraînement. Cela se formalise par l'échantillonnage à partir d'une postérieure biaisée par la récompense $q(x) \propto p_{\text{data}}(x)r(x)$.

Les méthodes actuelles de mise à l'échelle à l'inférence, en particulier le guidage, tentent d'atteindre cet objectif en modifiant la dérive de l'EDS d'échantillonnage avec un terme de gradient dérivé d'un potentiel de guidage $G$. Bien qu'efficaces empiriquement, les méthodes de guidage introduisent une limitation mathématique fondamentale : elles ne garantissent pas l'échantillonnage à partir de la distribution biaisée exacte $q(x)$. Un échantillonnage exact nécessite la transformation $h$ de Doob, qui implique la résolution d'une équation de Kolmogorov rétrograde et est souvent intraitable en haute dimension.

Les alternatives existantes, telles que FK-Corrector et AFDPS, tentent de corriger le biais de la génération guidée en utilisant le Monte Carlo Séquentiel (SMC) par rééchantillonnage des particules. Cependant, ces méthodes reposent sur des générateurs infinitésimaux qui nécessitent l'accès aux dérivées d'ordre supérieur de la fonction de récompense (spécifiquement le score $\nabla \log p_t$, le Laplacien $\Delta r$, et la Hessienne du guidage). Ces exigences entraînent une surcharge computationnelle significative et rendent ces méthodes inapplicables aux modèles de récompense boîte noire (par exemple, les modèles de récompense neuronaux utilisés dans la génération d'images à partir de texte) où les dérivées sont indisponibles ou coûteuses.

Le défi central abordé par ce travail est le suivant : Pouvons-nous concevoir une procédure de mise à l'échelle à l'inférence sans approximation qui élimine le biais du guidage sous-optimal sans nécessiter d'informations dérivées de la récompense ?

2. Méthodologie : URGE

Les auteurs proposent URGE (Unbiased Resampling via Girsanov Estimation), un algorithme de mise à l'échelle à l'inférence sans dérivée qui effectue un rééchantillonnage sur les trajectoires de génération (espace des trajectoires) plutôt que sur les échantillons générés individuellement (espace des particules).

Mécanisme Central

URGE exploite le théorème de Girsanov pour effectuer un changement de mesure sur l'espace des trajectoires du processus de diffusion.

1. Génération Guidée : Plusieurs trajectoires $\{X^{(i),G}_t\}$ sont simulées en utilisant une EDS guidée avec un potentiel de guidage arbitraire $G$ (qui peut être imparfait).
   $$dX^G_t = (v(X^G_t, t) + V(t)^2 \nabla_x G(X^G_t, t))dt + V(t)dW_t$$
2. Répondage par Trajectoire : Au lieu de calculer des poids basés sur des générateurs infinitésimaux (qui nécessitent des dérivées), URGE attribue un poids multiplicatif à chaque trajectoire entière basé sur la dérivée de Radon-Nikodym entre la mesure de trajectoire cible biaisée par la récompense $Q$ et la mesure de trajectoire guidée $P_G$.
   Le poids d'importance pour une trajectoire sur $[s, t]$ est :
   $$w^{\text{URGE}}_{s,t} = \exp\left( \int_s^t -V(\tau)\nabla_x G(X^G_\tau, \tau)^\top dW_\tau - \frac{1}{2}\int_s^t V(\tau)^2 \|\nabla_x G(X^G_\tau, \tau)\|^2 d\tau + \int_s^t dr(X^G_\tau, \tau) \right)$$
   Crucialement, cette formulation ne nécessite pas $\nabla r$, $\Delta r$, ou $\nabla \log p_t$. Elle requiert uniquement la fonction de récompense $r$ et le guidage $G$.
3. Rééchantillonnage : À des pas de temps discrets, les particules sont rééchantillonnées proportionnellement à leurs poids de trajectoire. Cela implémente un mécanisme de "naissance-mort" où les trajectoires à haute récompense sont répliquées et celles à faible récompense sont éliminées, conduisant la distribution empirique vers la postérieure vraie $q$.

Équivalence Théorique

L'article établit une équivalence théorique entre le rééchantillonnage dans l'espace des trajectoires (URGE) et le rééchantillonnage dans l'espace des particules (FK-Corrector/AFDPS).

• Les auteurs prouvent que le générateur marginalisé du processus d'espace de trajectoires URGE, lorsqu'il est conditionné par l'état terminal, retrouve le terme exact de rééquilibrage infinitésimal utilisé par FK-Corrector et AFDPS.
• Cela confirme que URGE est sans approximation : il converge vers la même distribution cible que les méthodes d'espace de particules théoriquement optimales (mais lourdes en dérivées), mais y parvient en relevant la construction des poids vers l'espace des trajectoires, contournant ainsi la nécessité de calculs explicites de dérivées.

3. Contributions Clés

1. Nouveau Cadre de Mélange d'Experts : Les auteurs introduisent une méthode qui intègre le SMC avec les processus génératifs en rééchantillonnant sur les trajectoires plutôt que sur les échantillons. Cela permet une construction flexible des poids via un changement de mesure de Girsanov sans nécessiter de dérivées de récompense.
2. Implémentation sans Dérivée : URGE élimine le besoin de fonctions de score, de Hessiennes ou de Laplaciens de la récompense. Cela le rend applicable aux modèles de récompense boîte noire (par exemple, CLIP, ImageReward, modèles ajustés par DPO) où les gradients sont inaccessibles.
3. Équivalence Théorique : L'article prouve que le rééquilibrage par trajectoire de URGE, lorsqu'il est marginalisé, retrouve exactement la structure de pondération des méthodes basées sur les dérivées les plus avancées (FK-Corrector/AFDPS), garantissant que les deux schémas produisent la même loi terminale sans approximation.
4. Supériorité Empirique : URGE surpasse systématiquement les références existantes (y compris FK-Steering et les variantes AFDPS) sur des tâches synthétiques et réelles, tout en étant significativement plus simple à mettre en œuvre et entièrement sans gradient.

4. Résultats Expérimentaux

Les auteurs évaluent URGE sur trois benchmarks :

• Modèle de Mélange Gaussien (Synthétique) :

  • URGE obtient l'erreur la plus faible sur toutes les métriques (MMD, SWD, Moyenne L2, Frob de Covariance) par rapport au Guidage Pur, AFDPS et FK-Steering.
  • Il démontre une efficacité de convergence supérieure, obtenant de meilleurs résultats avec moins de particules que les références.
  • URGE reste robuste face aux variations de dimensions et de pas de discrétisation, montrant une sensibilité moindre à la taille du pas que les méthodes concurrentes.

• Problèmes Inverses (Restauration d'Image) :

  • Testé sur le Défloutage Gaussien, le Défloutage par Mouvement, la Super-Résolution et l'Inpainting de Boîte utilisant ImageNet-256 et FFHQ-256.
  • URGE atteint des performances comparables à AFDPS (qui nécessite des dérivées) et surpasse significativement les autres références en PSNR et LPIPS.
  • URGE démontre une meilleure mise à l'échelle computationnelle, obtenant des temps d'inférence plus courts que AFDPS-ODE tout en maintenant une haute qualité de reconstruction.

• Génération d'Images à partir de Texte :

  • Testé sur Stable Diffusion v1.5 et SDXL avec des variantes ajustées par DPO.
  • URGE surpasse FK-Steering et le modèle de base en qualité esthétique (ImageReward, HPS) et en fidélité au prompt (GenEval, CLIP-Score).
  • Notamment, URGE permet à un modèle plus petit (SD v1.5) de correspondre ou de s'approcher des performances d'un modèle plus grand (SDXL) lors de l'utilisation de la mise à l'échelle par particules, alors que FK-Steering montre des rendements décroissants.
  • URGE présente une amélioration monotone de la récompense avec l'augmentation du nombre de particules, contrairement à FK-Steering.

5. Signification et Revendications

L'article revendique que URGE fournit un cadre d'échantillonnage de postérieure basé sur un ensemble, sans entraînement et sans approximation. Sa signification principale réside dans l'unification des stratégies de pondération basées sur les trajectoires et sur les états. En relevant le mécanisme de rééchantillonnage de l'espace des particules vers l'espace des trajectoires, URGE :

• Élimine le goulot d'étranglement des dérivées : Il permet l'utilisation de modèles de récompense complexes et boîte noire dans le guidage de diffusion sans le coût computationnel du calcul de dérivées d'ordre supérieur.
• Assure la correction théorique : Il garantit la convergence vers la vraie postérieure biaisée par la récompense, corrigeant le biais inhérent des méthodes de guidage standard.
• Offre une efficacité pratique : Il atteint une qualité de génération de pointe avec une implémentation plus simple et une surcharge computationnelle compétitive, le rendant adapté aux tâches génératives à grande échelle et réelles.

Les auteurs soulignent que, bien que les travaux précédents (FK-Corrector, AFDPS) aient fourni des insights théoriques sur la correction du guidage, ils étaient limités par leur dépendance aux informations dérivées. URGE comble cette lacune, offrant une solution pratique et sans dérivée qui conserve les garanties théoriques de l'échantillonnage sans approximation.