CMAD: Cooperative Multi-Agent Diffusion via Stochastic Optimal Control

Résumé Technique : CMAD – Diffusion Coopérative Multi-Agent via Contrôle Optimal Stochastique

Énoncé du Problème
Les modèles génératifs en temps continu, en particulier les modèles de diffusion basés sur le score, ont atteint des performances de pointe dans la synthèse et la restauration d'images. Bien que des techniques pour une génération contrôlable (par exemple, guidage par classifieur, modélisation conditionnelle du score) existent, un défi majeur subsiste dans la génération compositionnelle : agréger efficacement plusieurs modèles de diffusion pré-entraînés pour satisfaire des objectifs complexes et spécifiques à la tâche.

Les approches actuelles traitent largement la composition comme une combinaison algébrique de densités de probabilité (par exemple, produits d'experts ou mélanges). Ces méthodes supposent que la distribution cible est une fonction algébrique explicite des marginales composantes (par exemple, $p_{prod}(x) \propto \prod q_i(x)$). Cependant, les auteurs soutiennent que cette hypothèse est souvent irréaliste ; par exemple, il n'y a aucune garantie théorique que la combinaison d'un modèle pour le réalisme visuel avec un modèle pour l'alignement textuel aboutisse à une moyenne géométrique simple ou à un produit de densités. De plus, l'échantillonnage à partir de telles distributions de produits explicites nécessite souvent l'accès au score de la distribution combinée, qui ne se décompose pas linéairement, conduisant à des approximations biaisées et à une mauvaise qualité de génération.

Méthodologie
L'article propose CMAD (Cooperative Multi-Agent Diffusion), un changement de paradigme qui formule la génération compositionnelle comme un problème de Contrôle Optimal Stochastique Coopératif (SOC) plutôt que comme un problème de composition de densités.

1. Formulation Multi-Agent :

   • Les modèles de diffusion pré-entraînés sont traités comme des agents indépendants.
   • Chaque agent $i$ suit une Équation Différentielle Stochastique (EDS) contrôlée représentant sa dynamique en temps inverse :
     $$dX_t^{u,i} = [b_i(X_t^{u,i}, t) + g(t)u_i(X_t^{u,i}, t; \{X_t^{u,j}\}_j)]dt + g(t)dW_t$$
     Ici, $b_i$ est la dérive issue du modèle pré-entraîné, et $u_i$ est une politique de contrôle apprise. Crucialement, le contrôle $u_i$ dépend des états de tous les autres agents, créant un système d'EDS couplé.

2. Agrégation et Objectif :

   • L'objet généré final est défini comme un état agrégé $Y_t = \phi(\{X_t^{u,i}\}_{i=1}^N, t; \vartheta)$, où $\phi$ est un agrégateur (par exemple, masquage/concaténation non chevauchant) et $\vartheta$ sont des paramètres apprenables.
   • L'objectif est de minimiser une fonction objectif $J$ sur les contrôles $\{u_i\}$ et les paramètres de l'agrégateur $\vartheta$ :
     $$J = \mathbb{E}\left[ \int_0^T \left( \sum_{i=1}^N \lambda_i \|u_i\|^2 + c(Y_t, t) \right) dt + \Psi(Y_T) \right]$$
   • Coût Terminal ($\Psi$) : Encode la récompense spécifique à la tâche (par exemple, le logarithme négatif de la vraisemblance d'un classifieur pour un chiffre cible).
   • Coût de Fonctionnement ($c$) : Agit comme un stabilisateur et un substitut du coût terminal, évalué sur une "estimation de Tweedie" (une prévision débruitée) de l'état agrégé pour fournir des gradients denses le long de la trajectoire.
   • Coût de Contrôle : Le terme quadratique $\|u_i\|^2$ régularise le contrôle, maintenant les agents proches de leurs dynamiques en temps inverse pré-entraînées.

3. Schéma d'Optimisation (Descente par Contrôle) :

   • Les auteurs proposent un algorithme d'optimisation basé sur l'Optimisation Itérative de la Diffusion (IDO).
   • Au lieu de mettre à jour tous les contrôles conjointement (ce qui peut être coûteux en calcul), ils emploient un schéma de descente de coordonnées. À chaque itération externe, le contrôle d'un agent $u_i$ est mis à jour tandis que les autres sont figés.
   • La mise à jour pour un seul agent est effectuée via une simulation de Monte Carlo des trajectoires d'EDS couplées, en calculant les gradients de la perte empirique par rapport aux paramètres de contrôle, et en effectuant une étape de descente de gradient.
   • La politique de contrôle $u_i$ est paramétrée avec un biais inductif informé par la récompense, combinant une correction par réseau de neurones apprise avec un terme de guidage par gradient explicite dérivé de l'estimation de Tweedie de l'état agrégé.

Contributions Clés

1. Nouveau Cadre : Introduction d'un cadre multi-agent coopératif qui reformule l'inférence avec plusieurs modèles de diffusion pré-entraînés comme un problème de SOC, évitant ainsi le besoin de définitions algébriques explicites de la densité cible.
2. Algorithme d'Optimisation : Proposition d'un schéma d'optimisation par contrôle basé sur l'IDO, permettant un contrôle multi-agent évolutif en résolvant itérativement des sous-problèmes de SOC à agent unique.
3. Composition Implicite : Démonstration que des objectifs complexes (définis par des fonctions de perte) peuvent être optimisés sans connaître la forme de la densité composite résultante.

Résultats Expérimentaux
Le cadre a été validé sur une tâche de preuve de concept utilisant la génération de chiffres MNIST, où plusieurs agents étaient responsables de différentes bandes horizontales de l'image.

• Configuration : Les agents contrôlaient des régions non chevauchantes de l'image. Le coût terminal était le logarithme négatif de la vraisemblance d'un classifieur MNIST pré-entraîné pour un chiffre cible, augmenté d'une perte de continuité de suture pour assurer la cohérence spatiale.
• Lignes de Base : Comparaison avec une ligne de base naïve d'inférence (CDPS) utilisant un guidage par gradient sans contrôle coopératif appris.
• Performance :
  • CMAD (variantes conjointe et par contrôle) a systématiquement atteint une perte terminale ($\Psi$) plus faible par rapport à CDPS.
  • La précision de classification était élevée pour toutes les méthodes, mais CMAD a produit des échantillons plus réalistes visuellement.
  • CDPS produisait occasionnellement des artefacts non naturels aux suture entre les agents, tandis que CMAD générait des chiffres cohérents.
  • Le schéma d'optimisation par contrôle a atteint des performances compétitives (par exemple, 99,90 % de précision pour le chiffre 3 avec 3 agents) tout en offrant une alternative évolutive à l'optimisation conjointe.

Signification et Travaux Futurs
L'article revendique une importance en fournissant une alternative flexible à la composition basée sur les densités, permettant l'optimisation d'objectifs complexes et implicites définis par des fonctions de perte. Les auteurs notent que les résultats actuels sont des expériences initiales sur MNIST.

Les travaux futurs décrits dans l'article incluent :

• La mise à l'échelle du cadre vers des applications de haute dimension (par exemple, images naturelles).
• L'étude des conditions de convergence du schéma de mise à jour par contrôle.
• L'extension de la connexion aux jeux différentiels et aux dynamiques de jeu fictif.
• L'exploration de configurations où l'opérateur d'agrégation $\phi$ est lui-même paramétré et appris conjointement avec les contrôles.

Les auteurs maintiennent un ton modeste, présentant le travail comme une étape fondamentale vers la diffusion multi-agent coopérative plutôt que comme un benchmark résolu pour toutes les tâches génératives.