CMAD: Cooperative Multi-Agent Diffusion via Stochastic Optimal Control

Riepilogo Tecnico: CMAD – Diffusione Multi-Agente Cooperativa tramite Controllo Ottimo Stocastico

Enunciato del Problema
I modelli generativi a tempo continuo, in particolare i modelli di diffusione basati sul punteggio, hanno raggiunto prestazioni all'avanguardia nella sintesi e nel ripristino di immagini. Sebbene esistano tecniche per la generazione controllata (ad esempio, guida tramite classificatore, modellazione del punteggio condizionale), una sfida significativa permane nella generazione composizionale: aggregare efficacemente più modelli di diffusione pre-addestrati per soddisfare obiettivi complessi e specifici del compito.

Gli approcci attuali trattano prevalentemente la composizione come una combinazione algebrica di densità di probabilità (ad esempio, prodotti di esperti o miscele). Questi metodi assumono che la distribuzione target sia una funzione algebrica esplicita delle marginali dei componenti (ad esempio, $p_{prod}(x) \propto \prod q_i(x)$). Tuttavia, gli autori sostengono che tale assunzione è spesso irrealistica; ad esempio, non vi è alcuna garanzia teorica che combinare un modello per il realismo visivo con uno per l'allineamento al testo risulti in una semplice media geometrica o prodotto di densità. Inoltre, il campionamento da tali distribuzioni di prodotto esplicite richiede spesso l'accesso al punteggio della distribuzione combinata, che non si decompone linearmente, portando a approssimazioni distorte e a una scarsa qualità della generazione.

Metodologia
Il documento propone CMAD (Cooperative Multi-Agent Diffusion), un cambio di paradigma che formula la generazione composizionale come un problema di Controllo Ottimo Stocastico (SOC) Cooperativo piuttosto che come un problema di composizione di densità.

1. Formulazione Multi-Agente:

   • I modelli di diffusione pre-addestrati sono trattati come agenti indipendenti.
   • Ogni agente $i$ segue un'Equazione Differenziale Stocastica (SDE) controllata che rappresenta la sua dinamica a tempo inverso:
     $$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$$
     Qui, $b_i$ è la deriva del modello pre-addestrato e $u_i$ è una politica di controllo appresa. Crucialmente, il controllo $u_i$ dipende dagli stati di tutti gli altri agenti, creando un sistema di SDE accoppiate.

2. Aggregazione e Obiettivo:

   • L'oggetto generato finale è definito come uno stato aggregato $Y_t = \phi(\{X_t^{u,i}\}_{i=1}^N, t; \vartheta)$, dove $\phi$ è un aggregatore (ad esempio, mascheratura non sovrapposta/concatenazione) e $\vartheta$ sono parametri apprendibili.
   • L'obiettivo è minimizzare una funzione obiettivo $J$ sui controlli $\{u_i\}$ e sui parametri dell'aggregatore $\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]$$
   • Costo Terminale ($\Psi$): Codifica la ricompensa specifica del compito (ad esempio, il log-verosimiglianza negativo di un classificatore per una cifra target).
   • Costo di Esercizio ($c$): Agisce come stabilizzatore e surrogato del costo terminale, valutato su una "stima di Tweedie" (una previsione denoizzata) dello stato aggregato per fornire gradienti densi lungo la traiettoria.
   • Costo di Controllo: Il termine quadratico $\|u_i\|^2$ regolarizza il controllo, mantenendo gli agenti vicini alle loro dinamiche a tempo inverso pre-addestrate.

3. Schema di Ottimizzazione (Discesa per Controllo):

   • Gli autori propongono un algoritmo di ottimizzazione basato sull'Ottimizzazione Iterativa della Diffusione (IDO).
   • Invece di aggiornare tutti i controlli congiuntamente (il che può essere computazionalmente costoso), adottano uno schema di discesa per coordinate. In ogni iterazione esterna, il controllo $u_i$ di un agente viene aggiornato mentre gli altri vengono congelati.
   • L'aggiornamento per un singolo agente viene eseguito tramite simulazione Monte Carlo delle traiettorie delle SDE accoppiate, calcolando i gradienti della perdita empirica rispetto ai parametri di controllo ed eseguendo un passo di discesa del gradiente.
   • La politica di controllo $u_i$ è parametrizzata con un bias induttivo informato dalla ricompensa, combinando una correzione tramite rete neurale appresa con un termine esplicito di guida del gradiente derivato dalla stima di Tweedie dello stato aggregato.

Contributi Chiave

1. Nuovo Framework: Introduzione di un framework multi-agente cooperativo che inquadra l'inferenza con più modelli di diffusione pre-addestrati come un problema SOC, evitando la necessità di definizioni algebriche esplicite della densità target.
2. Algoritmo di Ottimizzazione: Proposta di uno schema di ottimizzazione per controllo basato su IDO, che abilita un controllo multi-agente scalabile risolvendo iterativamente sottoproblemi SOC per singolo agente.
3. Composizione Implicita: Dimostrazione che obiettivi complessi (definiti da funzioni di perdita) possono essere ottimizzati senza conoscere la forma della densità composita risultante.

Risultati Sperimentali
Il framework è stato validato su un compito di prova di concetto utilizzando la generazione di cifre MNIST, dove più agenti erano responsabili di diverse strisce orizzontali dell'immagine.

• Configurazione: Gli agenti controllavano regioni non sovrapposte dell'immagine. Il costo terminale era il log-verosimiglianza negativo di un classificatore MNIST pre-addestrato per una cifra target, arricchito con una perdita di continuità delle cuciture per garantire la coerenza spaziale.
• Baseline: Confronto con una baseline ingenua a tempo di inferenza (CDPS) che utilizza la guida del gradiente senza controllo cooperativo appreso.
• Prestazioni:
  • CMAD (sia le varianti congiunte che quelle per controllo) ha raggiunto costantemente una perdita terminale inferiore ($\Psi$) rispetto a CDPS.
  • L'accuratezza di classificazione è stata elevata per tutti i metodi, ma CMAD ha prodotto campioni visivamente più realistici.
  • CDPS ha occasionalmente prodotto artefatti innaturali alle cuciture tra gli agenti, mentre CMAD ha generato cifre coerenti.
  • Lo schema di ottimizzazione per controllo ha raggiunto prestazioni competitive (ad esempio, 99,90% di accuratezza per la cifra 3 con 3 agenti) offrendo al contempo un'alternativa scalabile all'ottimizzazione congiunta.

Significato e Lavori Futuri
Il documento rivendica importanza nel fornire un'alternativa flessibile alla composizione basata su densità, consentendo l'ottimizzazione di obiettivi complessi e impliciti definiti da funzioni di perdita. Gli autori notano che i risultati attuali sono esperimenti iniziali su MNIST.

I lavori futuri delineati nel documento includono:

• Scalare il framework ad applicazioni ad alta dimensionalità (ad esempio, immagini naturali).
• Investigare le condizioni di convergenza dello schema di aggiornamento per controllo.
• Estendere la connessione ai giochi differenziali e alle dinamiche di gioco fittizio.
• Esplorare contesti in cui l'operatore di aggregazione $\phi$ è esso stesso parametrizzato e appreso congiuntamente ai controlli.

Gli autori mantengono un tono modesto, presentando il lavoro come un passo fondazionale verso la diffusione multi-agente cooperativa piuttosto che come un benchmark risolto per tutti i compiti generativi.