CMAD: Cooperative Multi-Agent Diffusion via Stochastic Optimal Control

Resumo Técnico: CMAD – Difusão Cooperativa Multi-Agente via Controle Ótimo Estocástico

Declaração do Problema
Modelos generativos de tempo contínuo, particularmente modelos de difusão baseados em pontuação, alcançaram desempenho de última geração na síntese e restauração de imagens. Embora existam técnicas para geração controlada (por exemplo, orientação por classificador, modelagem de pontuação condicional), um desafio significativo permanece na geração composicional: agregar efetivamente múltiplos modelos de difusão pré-treinados para satisfazer objetivos complexos e específicos de tarefa.

As abordagens atuais tratam majoritariamente a composição como uma combinação algébrica de densidades de probabilidade (por exemplo, produtos de especialistas ou misturas). Esses métodos assumem que a distribuição-alvo é uma função algébrica explícita das marginais dos componentes (por exemplo, $p_{prod}(x) \propto \prod q_i(x)$). No entanto, os autores argumentam que essa suposição é frequentemente irrealista; por exemplo, não há garantia teórica de que combinar um modelo para realismo visual com um para alinhamento de texto resulte em uma média geométrica simples ou produto de densidades. Além disso, a amostragem de tais distribuições de produto explícitas frequentemente requer acesso à pontuação da distribuição combinada, a qual não se decompõe linearmente, levando a aproximações enviesadas e baixa qualidade de geração.

Metodologia
O artigo propõe o CMAD (Difusão Cooperativa Multi-Agente), uma mudança de paradigma que formula a geração composicional como um problema de Controle Ótimo Estocástico (SOC) Cooperativo, em vez de um problema de composição de densidades.

1. Formulação Multi-Agente:

   • Modelos de difusão pré-treinados são tratados como agentes independentes.
   • Cada agente $i$ segue uma Equação Diferencial Estocástica (EDE) controlada que representa sua dinâmica de tempo reverso:
     $$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$$
     Aqui, $b_i$ é a deriva do modelo pré-treinado, e $u_i$ é uma política de controle aprendida. Crucialmente, o controle $u_i$ depende dos estados de todos os outros agentes, criando um sistema de EDE acoplado.

2. Agregação e Objetivo:

   • O objeto gerado final é definido como um estado agregado $Y_t = \phi(\{X_t^{u,i}\}_{i=1}^N, t; \vartheta)$, onde $\phi$ é um agregador (por exemplo, mascaramento/concatenação não sobreposta) e $\vartheta$ são parâmetros aprendíveis.
   • O objetivo é minimizar uma função objetivo $J$ sobre os controles $\{u_i\}$ e os parâmetros do agregador $\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]$$
   • Custo Terminal ($\Psi$): Codifica a recompensa específica da tarefa (por exemplo, log-verossimilhança negativa de um classificador para um dígito alvo).
   • Custo de Execução ($c$): Atua como um estabilizador e substituto para o custo terminal, avaliado em uma "estimativa de Tweedie" (uma previsão desruidada) do estado agregado para fornecer gradientes densos ao longo da trajetória.
   • Custo de Controle: O termo quadrático $\|u_i\|^2$ regulariza o controle, mantendo os agentes próximos às suas dinâmicas de tempo reverso pré-treinadas.

3. Esquema de Otimização (Descida por Controle):

   • Os autores propõem um algoritmo de otimização baseado em Otimização Iterativa de Difusão (IDO).
   • Em vez de atualizar todos os controles conjuntamente (o que pode ser computacionalmente caro), eles empregam um esquema de descida coordenada. Em cada iteração externa, o controle de um agente $u_i$ é atualizado enquanto os outros são congelados.
   • A atualização para um único agente é realizada via simulação de Monte Carlo das trajetórias de EDE acopladas, computando gradientes da perda empírica em relação aos parâmetros de controle e executando um passo de descida de gradiente.
   • A política de controle $u_i$ é parametrizada com um viés indutivo informado por recompensa, combinando uma correção de rede neural aprendida com um termo explícito de orientação por gradiente derivado da estimativa de Tweedie do estado agregado.

Principais Contribuições

1. Novo Framework: Introdução de um framework cooperativo multi-agente que lança a inferência com múltiplos modelos de difusão pré-treinados como um problema de SOC, evitando a necessidade de definições algébricas explícitas da densidade-alvo.
2. Algoritmo de Otimização: Proposta de um esquema de otimização por controle baseado em IDO, permitindo controle multi-agente escalável resolvendo iterativamente subproblemas de SOC de agente único.
3. Composição Implícita: Demonstração de que objetivos complexos (definidos por funções de perda) podem ser otimizados sem conhecer a forma da densidade composta resultante.

Resultados Experimentais
O framework foi validado em uma tarefa de prova de conceito usando geração de dígitos MNIST, onde múltiplos agentes eram responsáveis por diferentes faixas horizontais da imagem.

• Configuração: Agentes controlavam regiões não sobrepostas da imagem. O custo terminal era a log-verossimilhança negativa de um classificador MNIST pré-treinado para um dígito alvo, aumentado com uma perda de continuidade de emenda para garantir coerência espacial.
• Baselines: Comparado contra uma baseline ingênua de tempo de inferência (CDPS) que usa orientação por gradiente sem controle cooperativo aprendido.
• Desempenho:
  • O CMAD (tanto variantes conjuntas quanto por controle) alcançou consistentemente menor perda terminal ($\Psi$) em comparação com o CDPS.
  • A precisão de classificação foi alta em todos os métodos, mas o CMAD produziu amostras mais visualmente realistas.
  • O CDPS ocasionalmente produziu artefatos não naturais nas emendas entre agentes, enquanto o CMAD gerou dígitos coerentes.
  • O esquema de otimização por controle alcançou desempenho competitivo (por exemplo, 99,90% de precisão para o dígito 3 com 3 agentes) enquanto oferecia uma alternativa escalável à otimização conjunta.

Significância e Trabalho Futuro
O artigo reivindica significância ao fornecer uma alternativa flexível à composição baseada em densidades, permitindo a otimização de objetivos complexos e implícitos definidos por funções de perda. Os autores observam que os resultados atuais são experimentos iniciais no MNIST.

O trabalho futuro delineado no artigo inclui:

• Escalar o framework para aplicações de alta dimensionalidade (por exemplo, imagens naturais).
• Investigar as condições de convergência do esquema de atualização por controle.
• Estender a conexão para jogos diferenciais e dinâmicas de jogo fictício.
• Explorar configurações onde o operador de agregação $\phi$ é ele próprio parametrizado e aprendido conjuntamente com os controles.

Os autores mantêm um tom modesto, apresentando o trabalho como um passo fundamental em direção à difusão cooperativa multi-agente, em vez de uma benchmark resolvida para todas as tarefas generativas.