Technical Summary: CMAD – Cooperative Multi-Agent Diffusion via Stochastic Optimal Control

Problem Statement

Continuous-time generative models, particularly score-based diffusion models and flow-based models, have achieved state-of-the-art performance in image restoration and synthesis. While significant progress has been made in controllable generation (e.g., via classifier guidance or conditional score modeling), a critical challenge remains: how to effectively aggregate multiple pre-trained diffusion models to satisfy complex, task-specific objectives.

Current compositional generation frameworks largely operate on the algebraic composition of probability densities (e.g., products of experts or mixtures). These methods assume the target distribution is an explicit algebraic combination of the marginals of the constituent models. The authors argue this assumption is often unrealistic; for instance, there is no theoretical guarantee that combining a model for visual realism with a model for text consistency results in a geometric average or a product of experts. Furthermore, these approaches require explicit knowledge of the resulting composite density at the diffusion-time level, which is rarely available for complex, implicit objectives defined by loss functions.

Methodology

The paper proposes CMAD (Cooperative Multi-Agent Diffusion), a paradigm shift that formulates compositional generation as a cooperative Stochastic Optimal Control (SOC) problem rather than a density composition problem.

Core Framework

1. Multi-Agent Formulation: Pre-trained diffusion models are treated as independent agents. Each agent $i$ follows a controlled Stochastic Differential Equation (SDE) representing the reverse-time diffusion process:
   $$dX^{u,i}_t = [b_i(X^{u,i}_t, t) + g(t)u_i(X^{u,i}_t, t; \{X^{u,j}_t\}_j)]dt + g(t)dW_t$$
   Here, $b_i$ is the drift of the pre-trained model, and $u_i$ is a learned control term. Crucially, the control $u_i$ depends on the states of all other agents, creating a coupled SDE system.

2. Aggregation and Objective: The generated object is defined as the aggregation of these interacting trajectories via an operator $\phi$:
   $$Y_t = \phi(\{X^{u,i}_t\}_{i=1}^N, t; \vartheta)$$
   The system optimizes a cooperative objective function $J$ over the controls $\{u_i\}$ and aggregation parameters $\vartheta$:
   $$J(\{u_i\}, \vartheta) = \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]$$

   • Control Cost: A weighted quadratic term penalizing deviations from the pre-trained reverse dynamics.
   • Running Cost ($c$): Encodes rewards along the trajectory, implemented as a surrogate of the terminal cost evaluated on a "Tweedie estimate" (a denoised look-ahead) of the aggregated state.
   • Terminal Cost ($\Psi$): A task-specific objective evaluated at the final time $T$ (e.g., negative log-likelihood of a classifier for a specific digit).

3. Optimization Scheme (Control-wise Descent):
   The paper introduces Iterative Diffusion Optimisation (IDO) adapted for multi-agent settings. Instead of updating all controls jointly, the framework employs a control-wise descent strategy (coordinate descent):

   • In each outer iteration, one agent's control $u_i$ is updated while others are frozen.
   • The update is performed via Monte Carlo simulation of the coupled SDE trajectories.
   • Gradients are computed using Backpropagation Through Time (BPTT) on the discretized SDE paths.
   • The control is parameterized with a reward-informed inductive bias, combining a learned neural network correction with an explicit gradient guidance term based on the Tweedie estimate.

Key Contributions

1. Cooperative Multi-Agent Framework: The authors introduce a novel framework that casts inference with multiple pre-trained diffusion models as a cooperative SOC problem, avoiding the need for explicit algebraic composition of densities.
2. Control-wise Optimization Scheme: They propose an iterative optimization scheme based on IDO that enables scalable multi-agent control by alternating updates between agents, effectively solving the coupled control problem.

Experimental Results

The framework was validated on a proof-of-concept task using MNIST digit generation.

• Setup: Multiple agents were assigned non-overlapping horizontal stripes of the image. The goal was to generate specific digits (0, 3, 9) by minimizing a terminal cost (classifier loss) and a seam-continuity loss (penalizing discontinuities between agent regions).
• Baselines: The method was compared against a naive inference-time baseline (CDPS) that uses gradient guidance without learned cooperative control.
• Performance:
  • Accuracy: CMAD achieved high classification accuracy across all digits and agent counts (2 and 3 agents), often surpassing the CDPS baseline (e.g., 99.90% vs 97.17% for digit 3 with 3 agents).
  • Terminal Loss: CMAD consistently attained lower terminal loss values compared to CDPS, indicating better optimization of the objective function.
  • Qualitative: Visual results showed that CMAD generated more realistic and coherent digit images compared to CDPS, which occasionally produced unnatural samples. However, the authors note that CMAD exhibited reduced diversity compared to the CDPS baseline.

Significance and Claims

The paper claims that CMAD offers a pragmatic and flexible alternative to density-based composition. By treating composition as a control problem, the framework can optimize for complex objectives defined by loss functions without requiring explicit knowledge of the form of the resulting composite density.

The authors position this work as a foundational step toward scalable, high-dimensional compositional generation. They acknowledge that the current results are limited to MNIST and that future work must address scaling to higher dimensions, potentially requiring extensions to path-wise gradient estimators or adjoint matching techniques. The paper modestly frames the current contribution as a demonstration of the framework's suitability for compositional tasks, with future investigations planned into convergence conditions, connections to differential games, and learning the aggregation operator jointly with the controls.