CMAD: Cooperative Multi-Agent Diffusion via Stochastic Optimal Control

기술적 요약: CMAD – 확률적 최적 제어를 통한 협력적 다중 에이전트 확산

문제 제기
연속 시간 생성 모델, 특히 스코어 기반 확산 모델은 이미지 합성 및 복원에서 최첨단 성능을 달성했습니다. 제어 가능한 생성 (예: 분류기 안내, 조건부 스코어 모델링) 을 위한 기술은 존재하지만, 구성적 생성 (compositional generation): 복잡한 작업별 목표를 충족시키기 위해 여러 사전 훈련된 확산 모델을 효과적으로 집계하는 것은 여전히 중요한 과제로 남아 있습니다.

현재의 접근 방식은 대부분 확률 밀도의 대수적 결합 (예: 전문가의 곱이나 혼합) 으로 구성을 다룹니다. 이러한 방법들은 목표 분포가 구성 요소 주변 분포의 명시적 대수 함수 (예: $p_{prod}(x) \propto \prod q_i(x)$) 라고 가정합니다. 그러나 저자들은 이 가정이 종종 비현실적이라고 주장합니다. 예를 들어, 시각적 사실감을 위한 모델과 텍스트 정렬을 위한 모델을 결합한다고 해서 밀도의 단순한 기하학적 평균이나 곱이 된다는 이론적 보장은 없습니다. 또한, 이러한 명시적 곱 분포로부터 샘플링하는 것은 종종 결합 분포의 스코어에 대한 접근을 필요로 하는데, 이는 선형적으로 분해되지 않아 편향된 근사와 낮은 생성 품질로 이어집니다.

방법론
이 논문은 **CMAD(협력적 다중 에이전트 확산)**를 제안하며, 이는 구성적 생성을 밀도 결합 문제가 아닌 협력적 확률적 최적 제어 (SOC) 문제로 공식화하는 패러다임의 전환을 의미합니다.

1. 다중 에이전트 공식화:

   • 사전 훈련된 확산 모델은 독립적인 에이전트로 간주됩니다.
   • 각 에이전트 $i$는 역시간 역학을 나타내는 제어된 확률 미분 방정식 (SDE) 을 따릅니다:
     $$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$$
     여기서 $b_i$는 사전 훈련된 모델에서의 드리프트이며, $u_i$는 학습된 제어 정책입니다. 결정적으로, 제어 $u_i$는 모든 다른 에이전트의 상태에 의존하여 결합된 SDE 시스템을 생성합니다.

2. 집계 및 목적:

   • 최종 생성된 객체는 집계된 상태 $Y_t = \phi(\{X_t^{u,i}\}_{i=1}^N, t; \vartheta)$로 정의되며, 여기서 $\phi$는 집계기 (예: 겹치지 않는 마스킹/연결) 이고 $\vartheta$는 학습 가능한 매개변수입니다.
   • 목표는 제어 $\{u_i\}$와 집계기 매개변수 $\vartheta$에 대한 목적 함수 $J$를 최소화하는 것입니다:
     $$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]$$
   • 종단 비용 ($\Psi$): 작업별 보상을 인코딩합니다 (예: 목표 숫자에 대한 분류기의 음의 로그 가능도).
   • 실행 비용 ($c$): 안정화제이자 종단 비용의 대리 역할을 하며, 궤적을 따라 밀집된 기울기를 제공하기 위해 집계된 상태의 "트위디 (Tweedie) 추정치"(탈노이즈 미리보기) 에 대해 평가됩니다.
   • 제어 비용: 2 차 항 $\|u_i\|^2$는 제어를 정규화하여 에이전트가 사전 훈련된 역시간 역학에 가깝게 유지되도록 합니다.

3. 최적화 체계 (제어별 하강):

   • 저자들은 **반복적 확산 최적화 (IDO)**에 기반한 최적화 알고리즘을 제안합니다.
   • 모든 제어를 동시에 업데이트하는 것 (계산 비용이 많이 들 수 있음) 대신, 좌표 하강 (coordinate-descent) 체계를 사용합니다. 각 외부 반복에서 다른 에이전트는 고정된 채로 한 에이전트의 제어 $u_i$만 업데이트됩니다.
   • 단일 에이전트에 대한 업데이트는 결합된 SDE 궤적의 몬테카를로 시뮬레이션을 통해 수행되며, 제어 매개변수에 대한 경험적 손실의 기울기를 계산한 후 기울기 하강 단계를 수행합니다.
   • 제어 정책 $u_i$는 보상 기반 귀납적 편향으로 매개변수화되며, 이는 집계된 상태의 트위디 추정치에서 유도된 명시적 기울기 안내 항과 결합된 학습된 신경망 보정을 포함합니다.

주요 기여

1. 새로운 프레임워크: 목표 밀도의 명시적 대수적 정의가 필요하지 않도록, 여러 사전 훈련된 확산 모델로 추론하는 것을 SOC 문제로 제시하는 협력적 다중 에이전트 프레임워크의 도입.
2. 최적화 알고리즘: 단일 에이전트 SOC 부분 문제를 반복적으로 해결함으로써 확장 가능한 다중 에이전트 제어를 가능하게 하는 IDO 기반의 제어별 최적화 체계 제안.
3. 암시적 결합: 손실 함수로 정의된 복잡한 목표를 결과 결합 밀도의 형태를 알지 않고도 최적화할 수 있음을 입증.

실험 결과
이 프레임워크는 MNIST 숫자 생성을 위한 개념 증명 작업에서 검증되었으며, 여기서 여러 에이전트가 이미지의 서로 다른 가로 줄무늬를 담당했습니다.

• 설정: 에이전트들은 이미지의 겹치지 않는 영역을 제어했습니다. 종단 비용은 목표 숫자에 대한 사전 훈련된 MNIST 분류기의 음의 로그 가능도였으며, 공간적 일관성을 보장하기 위해 이음새 연속성 손실이 추가되었습니다.
• 기준선: 학습된 협력적 제어 없이 기울기 안내를 사용하는 단순한 추론 시간 기준선 (CDPS) 과 비교했습니다.
• 성능:
  • CMAD(연결 및 제어별 변형 모두) 는 CDPS 에 비해 일관되게 더 낮은 종단 손실 ($\Psi$) 을 달성했습니다.
  • 분류 정확도는 모든 방법에서 높았지만, CMAD 는 더 시각적으로 사실적인 샘플을 생성했습니다.
  • CDPS 는 에이전트 간의 이음새에서 때때로 비자연적인 아티팩트를 생성한 반면, CMAD 는 일관된 숫자를 생성했습니다.
  • 제어별 최적화 체계는 경쟁력 있는 성능을 달성했습니다 (예: 3 에이전트로 숫자 3 의 정확도 99.90%) 동시 최적화에 대한 확장 가능한 대안을 제공했습니다.

의의 및 향후 작업
이 논문은 밀도 기반 결합에 대한 유연한 대안을 제공함으로써 손실 함수로 정의된 복잡하고 암시적인 목표의 최적화를 가능하게 한다는 점에서 의의가 있다고 주장합니다. 저자들은 현재 결과가 MNIST 에 대한 초기 실험임을 지적합니다.

논문에 명시된 향후 작업에는 다음이 포함됩니다:

• 고차원 응용 분야(예: 자연 이미지) 로 프레임워크 확장.
• 제어별 업데이트 체계의 수렴 조건 조사.
• 미분 게임 및 허구적 플레이 역학과의 연결 확장.
• 집계 연산자 $\phi$ 자체가 매개변수화되어 제어와 함께 학습되는 설정 탐구.

저자들은 모든 생성 작업에 대한 해결된 벤치마크가 아니라 협력적 다중 에이전트 확산을 향한 기초적인 단계로 이 작업을 제시하며 겸손한 어조를 유지합니다.