Résumé Technique : Modèles de Diffusion Contraints avec Inférence Primal-Dual

Énoncé du Problème

Les modèles de diffusion standards excellent dans l'échantillonnage de distributions complexes à haute dimension dans des contextes non contraints (par exemple, la synthèse d'images et de sons). Cependant, ils manquent de mécanismes robustes pour imposer des contraintes statistiques (moyennes) lors du processus d'échantillonnage. De nombreux problèmes d'optimisation du monde réel — tels que l'allocation de ressources sans fil ou la gestion de portefeuille — nécessitent des solutions qui sont probabilistes plutôt que déterministes, où la politique optimale est une distribution satisfaisant des contraintes en espérance (par exemple, des garanties de débit par utilisateur ou des budgets de risque).

Les approches existantes pour l'échantillonnage contraint reposent souvent sur :

1. Contraintes ponctuelles : Imposer des contraintes sur chaque échantillon individuel (par exemple, via des projections ou des réflexions de bord), ce qui peut être excessivement conservateur et dégrader l'optimalité.
2. Entraînement Dual (DT - Dual Training) : Estimer le multiplicateur de Lagrange optimal (variable duale) avant l'échantillonnage et le figer. Cette approche est fragile car le multiplicateur optimal est défini par des espérances intraitables sous la distribution même qui est apprise, nécessant un réentraînement coûteux pour chaque nouvelle instance de problème ou changement de contrainte. De plus, un échantillonneur figé ne peut pas s'adapter aux violations de contraintes observées pendant la génération.

Méthodologie : Inférence Primal-Dual (PDI)

Le papier propose l'Inférence Primal-Dual (PDI), un cadre qui traite la variable duale non pas comme un paramètre pré-calculé, mais comme un état dynamique de l'échantillonneur qui évolue conjointement avec le processus de diffusion inverse.

1. Formulation Théorique

Le problème est formulé comme une optimisation régularisée par l'entropie dans le domaine dual de Lagrange. La distribution optimale $\mu^*$ est montrée comme étant une distribution de Gibbs indexée par la variable duale optimale $\lambda^*$ :
$$\mu^*(x) \propto \exp\left(-\frac{1}{\beta}(f_0(x) + \lambda^{*\top}f(x))\right)$$
où $f_0$ est l'objectif, $f$ représente les contraintes, et $\beta$ est la température.

2. L'Algorithme PDI

PDI entrelace la dénonisation primale (étapes de diffusion inverse) avec l'ascension duale (mise à jour du multiplicateur) à chaque étape d'inférence $t$ :

• Étape Primale (PDI-P) : L'échantillonneur met à jour l'échantillon bruité $x_t$ en utilisant un champ de score conditionné par la variable duale actuelle $\lambda_t$.
• Étape Duale (PDI-D) : La variable duale est mise à jour via l'ascension de gradient basée sur l'estimation de la violation de contrainte des échantillons dénoisés. Crucialement, la violation de contrainte est estimée en utilisant la formule de Tweedie (estimations de la moyenne a posteriori) plutôt que des échantillons propres, permettant ainsi la mise à jour sans accès à la vérité terrain.
  $$\lambda_{t+1} = \left[ \lambda_t + \eta_t \mathbb{E}_{x_{t+1}} \left[ f(\hat{y}_0) \right] \right]_+$$
  où $\hat{y}_0$ est l'estimation de Tweedie de l'échantillon propre.

3. Réseau de Score Entraîné

Pour permettre ce couplage dynamique, les auteurs entraînent un réseau de score unique conditionné par la dualité $s_\theta(x_t, t, \lambda_t)$. Contraなる aux modèles standards entraînés sur une seule distribution cible, ce réseau apprend la famille de champs de score de Gibbs indexés par la plage de variables duales rencontrées pendant la trajectoire d'inférence. La procédure d'entraînement combine :

• Exploitation : Déroulement du modèle actuel pour collecter des trajectoires de paires $(x_t, \lambda_t)$.
• Exploration : Échantillonnage de variables duales à partir d'une distribution a priori pour couvrir l'espace des multiplicateurs potentiels.
• Estimation de Monte Carlo : Utilisation d'estimateurs MC pour approximer le véritable score de la distribution de Gibbs pour les cibles de régression.

Principales Contributions

Le papier présente cinq contributions principales :

1. Formulation de Point-Selle : Formuler l'échantillonnage contraint avec des contraintes moyennes comme un problème de point-selle dans le domaine dual de Lagrange, identifiant la distribution optimale comme une distribution de Gibbs indexée par la variable duale.
2. Algorithme PDI : Introduction d'un échantillonneur de diffusion inverse où le dénoising primal et l'ascension duale évoluent conjointement. Le multiplicateur est mis à jour à chaque étape sur la base des violations de contraintes estimées à partir des échantillons dénoisés.
3. Réseau de Score Généralisé : Entraînement d'un réseau unique conditionné par la variable duale, permettant au modèle de généraliser à travers une famille de distributions de Gibbs plutôt qu'à une seule cible fixe.
4. Analyse de Convergence : Prouver que les itérés duals moyennés dans le temps convergent vers un voisinage du multiplicateur optimal à un taux de $O(1/\sqrt{T})$. Les auteurs établissent également une borne sur l'effet du déséquilibre dual résiduel sur la distribution terminale, montrant que les étapes initiales à haut bruit atténuent les erreurs tandis que les étapes ultérieures affinent la solution.
5. Validation Empirique : Démontrer l'efficacité de la méthode sur trois problèmes d'optimisation stochastique distincts : l'échantillonnage de mélanges gaussiens, le contrôle de puissance sans fil et la gestion de portefeuille.

Résultats Expérimentaux

Les auteurs évaluent PDI par rapport à des bases incluant les modèles de diffusion projetés (PDM), le Langevin Primal-Dual (PDL), l'échantillonnage de postérieur par diffusion (DPS) et l'entraînement dual (DT).

• Mélange de Gaussiennes : PDI atteint de meilleures valeurs d'objectif (plus proches de l'optimum non contraint) tout en maintenant une pleine faisabilité des contraintes. Contrairement à PDM et PDL, qui imposent des contraintes sur chaque échantillon (menant à des solutions conservatrices), PDI satisfait les contraintes en espérance, permettant des échantillons plus diversifiés et optimaux.
• Allocation de Ressources Sans Fil : PDI atteint le meilleur équilibre entre le débit moyen (optimalité) et la faisabilité (satisfaction des contraintes de débit minimum pour les utilisateurs). Il surpasse significativement PDM et PDL dans les métriques de faisabilité (ex: débits au 5ème percentile) tout en maintenant des débits moyens élevés. Il démontre également une robustesse aux changements de contraintes hors distribution (OOD) (ex: différentes valeurs de $r_{min}$) par rapport à DT.
• Gestion de Portefeuille : PDI fournit le meilleur compromis entre le rendement attendu, la faisabilité du risque et la diversification du portefeuille (entropie). Il surpasse les modèles non contraints (qui violent les budgets de risque) et les bases trop conservatrices (qui sacrifient le rendement).

Études d'Ablation :

• Multiplicateurs Fixes : Figer la variable duale (soit l'itéré final, soit la moyenne temporelle) durant l'inférence entraîne des performances inférieures à la dynamique complète de PDI, confirmant que l'évolution dépendante du chemin de la variable duale est critique.
• Planning de Bruit (Noise Schedule) : Le choix du planning de bruit est vital. Les schémas qui maintiennent un rapport signal sur bruit (SNR) faible plus longtemps (ex: cosinus standard) permettent à la variable duale de se stabiliser avant que le processus inverse ne devienne sensible aux erreurs. Les schémas qui reconstruisent les échantillons propres trop tôt amplifient les déséquilibres duaux, menant à des violations de contraintes.
• Entraînement Dual vs PDI : Bien que DT puisse atteindre des variables duales informatives, il ne parvient pas à égaler la performance de PDI car son modèle de score est un champ "moyenné" sur la trajectoire d'entraînement et ne peut pas s'adapter aux résidus de contrainte spécifiques durant l'inférence.

Signification et Revendications

Le papier affirme que PDI comble le fossé entre la modélisation générative basée sur la diffusion et l'optimisation stochastique contrainte. Sa signification réside dans :

• Adaptation Dynamique : Contrairement aux méthodes antérieures qui nécessitent de pré-calculer et de figer le multiplicateur dual, PDI traite le multiplicateur comme un état d'inférence, permettant à l'échantillonneur de réagir instantanément aux violations de contraintes.
• Efficacité et Robustesse : En entraînant un réseau unique sur une famille de distributions, PDI évite le réentraînement coûteux requis par l'entraînement dual pour de nouvelles instances de problèmes ou des changements de contraintes.
• Garanties Théoriques : Ce travail fournit des garanties de convergence pour les variables duales moyennées dans le temps et des bornes de stabilité pour le processus inverse, offrant une base théorique pour l'échantillonnage de diffusion contraint.

Les auteurs notent des limites, spécifiquement que les garanties de convergence s'appliquent aux itérés duals moyennés dans le temps plutôt qu'à la distribution primale finale directement, et que la méthode nécessite un réglage des hyperparamètres spécifique au problème (ex: force de régularisation de l'entropie $\beta$ et priors duaux). Néanmoins, PDI est présenté comme un cadre robuste pour l'échantillonnage de distributions optimales dans des problèmes d'optimisation contraints où les solutions doivent être probabilistes.