CMAD: Cooperative Multi-Agent Diffusion via Stochastic Optimal Control

Technische Samenvatting: CMAD – Cooperatieve Multi-Agent Diffusie via Stochastische Optimale Controle

Probleemstelling
Continue-tijd generatieve modellen, met name score-gebaseerde diffusiemodellen, hebben state-of-the-art prestaties behaald in beeldsynthese en -herstel. Hoewel er technieken bestaan voor controleerbare generatie (bijvoorbeeld classifier-guidance, conditionele score-modellering), blijft een aanzienlijke uitdaging bestaan bij compositional generation: het effectief aggregeren van meerdere vooraf getrainde diffusiemodellen om complexe, taakspecifieke doelen te vervullen.

Huidige benaderingen behandelen compositie grotendeels als een algebraïsche combinatie van kansdichtheden (bijvoorbeeld producten van experts of mengsels). Deze methoden gaan ervan uit dat de doeldistributie een expliciete algebraïsche functie is van de componentmarginalen (bijvoorbeeld $p_{prod}(x) \propto \prod q_i(x)$). De auteurs betogen echter dat deze aanname vaak onrealistisch is; er is bijvoorbeeld geen theoretisch bewijs dat het combineren van een model voor visuele realisme met een model voor tekstuitlijning resulteert in een eenvoudig geometrisch gemiddelde of product van dichtheden. Bovendien vereist het stalen uit dergelijke expliciete productdistributies vaak toegang tot de score van de gecombineerde distributie, die niet lineair decomposeert, wat leidt tot bevooroordeelde benaderingen en slechte generatiekwaliteit.

Methodologie
Het artikel stelt CMAD (Cooperative Multi-Agent Diffusion) voor, een paradigmaverschuiving die compositional generation formuleert als een Coöperatief Stochastisch Optimale Controle (SOC)-probleem in plaats van een dichtheidscompositieprobleem.

1. Multi-Agent Formulering:

   • Vooraf getrainde diffusiemodellen worden behandeld als onafhankelijke agenten.
   • Elke agent $i$ volgt een gecontroleerde Stochastische Differentiaalvergelijking (SDE) die zijn dynamica in omgekeerde tijd vertegenwoordigt:
     $$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$$
     Hierbij is $b_i$ de drift van het vooraf getrainde model, en $u_i$ een geleerde controlebeleid. Cruciaal is dat de controle $u_i$ afhangt van de toestanden van alle andere agenten, waardoor een gekoppeld SDE-systeem ontstaat.

2. Aggregatie en Doel:

   • Het gegenereerde eindobject wordt gedefinieerd als een geaggregeerde toestand $Y_t = \phi(\{X_t^{u,i}\}_{i=1}^N, t; \vartheta)$, waarbij $\phi$ een aggregator is (bijvoorbeeld niet-overlappend maskeren/concatenatie) en $\vartheta$ leerbare parameters zijn.
   • Het doel is om een doelfunctie $J$ te minimaliseren over de controles $\{u_i\}$ en aggregatorparameters $\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]$$
   • Eindkosten ($\Psi$): Codeert de taakspecifieke beloning (bijvoorbeeld de negatieve log-likelihood van een classifier voor een doelcijfer).
   • Doorlopende Kosten ($c$): Dient als stabilisator en vervanger voor de eindkosten, geëvalueerd op een "Tweedie-schatting" (een gedenoised vooruitblik) van de geaggregeerde toestand om dichte gradiënten langs de trajecten te leveren.
   • Controlekosten: De kwadratische term $\|u_i\|^2$ regulariseert de controle, waardoor de agenten dicht bij hun vooraf getrainde dynamica in omgekeerde tijd blijven.

3. Optimalisatieschema (Controle-gerichte Afdaling):

   • De auteurs stellen een optimalisatiealgoritme voor gebaseerd op Iterative Diffusion Optimisation (IDO).
   • In plaats van alle controles gezamenlijk te updaten (wat rekenkundig duur kan zijn), hanteren ze een coördinaat-afdalingsschema. In elke buitenste iteratie wordt de controle $u_i$ van één agent bijgewerkt terwijl de anderen bevroren blijven.
   • De update voor een enkele agent wordt uitgevoerd via Monte Carlo-simulatie van de gekoppelde SDE-trajecten, het berekenen van gradiënten van de empirische verliezen met betrekking tot de controleparameters, en het uitvoeren van een gradiëntafdalingstap.
   • Het controlebeleid $u_i$ wordt geparametriseerd met een beloningsgeïnformeerde inductieve bias, waarbij een geleerde neurale netwerkcorrectie wordt gecombineerd met een expliciete gradiëntgeleidingsterm afgeleid van de Tweedie-schatting van de geaggregeerde toestand.

Belangrijkste Bijdragen

1. Nieuw Kader: Introductie van een coöperatief multi-agentkader dat inferentie met meerdere vooraf getrainde diffusiemodellen omzet in een SOC-probleem, waardoor de noodzaak voor expliciete algebraïsche definities van de doeldichtheid wordt vermeden.
2. Optimalisatiealgoritme: Voorstellen van een controle-gerichte optimalisatieschema gebaseerd op IDO, wat schaalbare multi-agentcontrole mogelijk maakt door iteratief SOC-deelproblemen voor enkele agenten op te lossen.
3. Impliciete Compositie: Aantonen dat complexe doelen (gedefinieerd door verliesfuncties) kunnen worden geoptimaliseerd zonder de vorm van de resulterende compositiedichtheid te kennen.

Experimentele Resultaten
Het kader werd gevalideerd op een proof-of-concept-taak met MNIST-cijfergeneratie, waarbij meerdere agenten verantwoordelijk waren voor verschillende horizontale strepen van het beeld.

• Opstelling: Agenten controleerden niet-overlappende regio's van het beeld. De eindkosten waren de negatieve log-likelihood van een vooraf getrainde MNIST-classifier voor een doelcijfer, aangevuld met een naad-continuïteitsverlies om ruimtelijke coherentie te waarborgen.
• Basislijnen: Vergelijking met een naïeve inferentie-tijdbasislijn (CDPS) die gradiëntgeleiding gebruikt zonder geleerde coöperatieve controle.
• Prestaties:
  • CMAD (zowel gezamenlijke als controle-gerichte varianten) behaalde consistent lagere eindkosten ($\Psi$) in vergelijking met CDPS.
  • De classificatie-accuraatheid was hoog voor alle methoden, maar CMAD produceerde visueel realistischere samples.
  • CDPS produceerde af en toe onnatuurlijke artefacten bij de naden tussen agenten, terwijl CMAD coherente cijfers genereerde.
  • Het controle-gerichte optimalisatieschema behaalde concurrerende prestaties (bijvoorbeeld 99,90% accurate voor cijfer 3 met 3 agenten) terwijl het een schaalbaar alternatief bood voor gezamenlijke optimalisatie.

Betekenis en Toekomstig Werk
Het artikel claimt betekenis door een flexibel alternatief te bieden voor op dichtheid gebaseerde compositie, waardoor optimalisatie mogelijk is van complexe, impliciete doelen gedefinieerd door verliesfuncties. De auteurs merken op dat de huidige resultaten initiële experimenten zijn op MNIST.

Toekomstig werk zoals beschreven in het artikel omvat:

• Schalen van het kader naar hoogdimensionale toepassingen (bijvoorbeeld natuurlijke afbeeldingen).
• Onderzoeken van de convergentievoorwaarden van het controle-gerichte updateschema.
• Uitbreiden van de connectie met differentiële spellen en fictieve-spel-dynamica.
• Verkennen van scenario's waarbij de aggregatoroperator $\phi$ zelf geparametriseerd is en gezamenlijk met de controles wordt geleerd.

De auteurs handhaven een bescheiden toon en presenteren het werk als een fundamentele stap naar coöperatieve multi-agent diffusie in plaats van een opgeloste benchmark voor alle generatieve taken.