MASPOB: Bandit-Based Prompt Optimization for Multi-Agent Systems with Graph Neural Networks

이 논문은 다중 에이전트 시스템의 프롬프트 최적화에서 샘플 효율성, 토폴로지 결합, 그리고 검색 공간의 복잡성이라는 세 가지 주요 과제를 해결하기 위해 밴딧 알고리즘과 그래프 신경망 (GNN) 을 결합한 새로운 프레임워크인 MASPOB 를 제안하고, 이를 통해 기존 베이스라인을 능가하는 최첨단 성능을 입증합니다.

핵심

🍽️ 비유: 고급 레스토랑의 요리 팀

상상해 보세요. 거대한 요리를 완성하기 위해 주방장 (Agent 1), 소스 담당 (Agent 2), 디저트 담당 (Agent 3) 등 여러 명의 요리사들이 팀을 이루고 있습니다.

1. 문제 상황 (기존 방식의 한계):

   • 이 팀은 이미 완벽하게 설계된 '레시피 (워크플로우)'를 따르고 있습니다. 주방장에게 소스를 만들라고 하거나, 소스 담당에게 메인 요리를 하라고 지시하는 것은 변경할 수 없습니다. (실제 산업 현장에서는 안전이나 규정 때문에 이런 구조를 바꾸는 건 불가능한 경우가 많습니다.)
   • 하지만 각 요리사가 **"어떤 말 (프롬프트) 을 듣고 시작하느냐"**에 따라 요리의 맛은 천차만별입니다.
   • 문제는, 이 팀의 실력을 테스트하려면 매번 실제 요리를 다 만들어서 맛을 봐야 한다는 점입니다. 이 과정은 시간과 돈이 너무 많이 듭니다. (논문에서 말하는 '높은 평가 비용')
   • 게다가, 주방장의 말이 바뀌면 소스 담당이 받는 재료가 달라지고, 그 결과 디저트까지 영향을 받습니다. 서로가 서로에게 영향을 미쳐서 (Topology-induced coupling) 따로따로 연습하는 것만으로는 최고의 맛을 낼 수 없습니다.

2. 기존 방법들의 부족함:

   • 기존 기술들은 각 요리사에게 "더 맛있게 해봐"라고 막연히 말하거나, 무작위로 레시피를 바꿔보며 시도해봤습니다. 하지만 서로의 영향을 고려하지 않아서 비효율적이었습니다.

3. MASPOB 의 해결책 (이 기술의 핵심):
   MASPOB 는 이 문제를 해결하기 위해 세 가지 마법 같은 도구를 사용합니다.

   • ① '운세 점수'와 '확실한 맛'의 균형 (Bandit & UCB):

     • 모든 레시피를 다 만들어 볼 돈이 없으니, 가장 유망한 레시피를 고르는 동시에 아직 시도해본 적 없는 새로운 레시피도 가끔 시도해봐야 합니다.
     • MASPOB 는 마치 "이 레시피는 확실히 맛있을 것 같아 (Exploitation)"와 "이건 아직 안 해봤는데, 혹시 대박일지도 몰라 (Exploration)"를 동시에 계산하는 스마트한 점수판을 사용합니다. 이 점수판을 통해 적은 시도 횟수로 최고의 조합을 찾아냅니다.

   • ② '팀의 관계도'를 읽는 눈 (Graph Neural Network, GNN):

     • 이 기술은 각 요리사 (Agent) 가 서로 어떻게 연결되어 있는지 **관계도 (그래프)**를 그립니다.
     • "주방장이 A 를 하면 소스 담당은 B 를 해야 하고, 그다음 디저트 담당은 C 를 해야 해"라는 팀의 구조를 AI 가 미리 학습합니다.
     • 그래서 "주방장의 말을 조금만 바꿔도 소스 담당의 반응이 이렇게 변할 거야"라고 예측할 수 있게 되어, 실제 요리를 다 만들어보지 않아도 어떤 조합이 좋을지 감을 잡습니다.

   • ③ '한 번에 하나씩' 수정하기 (Coordinate Ascent):

     • 모든 요리사의 말을 한 번에 다 바꿔보려면 조합의 수가 너무 많아져서 (수백만 가지) 불가능합니다.
     • MASPOB 는 한 명씩 말을 바꿔가며 최적의 조합을 찾습니다. "오늘은 주방장만 말을 바꿔보자. 그다음 소스 담당만 바꿔보자." 이렇게 순서대로 하나씩 고쳐가면서 전체적인 맛을 끌어올립니다.

🏆 결과: 무엇이 달라졌나요?

이 기술을 적용한 결과, 기존 방법들보다 **훨씬 적은 비용 (적은 시도 횟수)**으로 **더 맛있는 요리 (높은 성능)**를 만들어냈습니다.

• 질문 답변 (HotpotQA), 코드 작성 (HumanEval), 수학 문제 (GSM8K) 등 다양한 분야에서 기존 최고의 방법들보다 더 좋은 성적을 거두었습니다.
• 특히, 팀의 구조를 바꾸지 않고도 각자의 역할 (프롬프트) 을 잘 다듬는 것만으로도 팀 전체의 실력이 비약적으로 상승함을 증명했습니다.

💡 한 줄 요약

**"이미 완성된 팀 구조를 바꾸지 않고, 각 팀원이 서로의 영향을 고려하며 '최고의 말'을 찾아내는 스마트한 코치 시스템"**입니다.

이 기술은 AI 가 복잡한 일을 할 때, 막대한 비용 없이도 효율적으로 최고의 성과를 낼 수 있게 해주는 현실적인 해결책을 제시합니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem Setting)

대형 언어 모델 (LLM) 기반의 다중 에이전트 시스템 (MAS, Multi-Agent Systems) 은 복잡한 작업을 해결하기 위해 여러 에이전트가 워크플로우를 통해 협력하는 방식으로 널리 사용되고 있습니다. 그러나 실제 배포 환경 (의료, 금융 감사 등) 에서는 안전성 검증 및 규정 준수를 위해 워크플로우 토폴로지 (구조) 를 변경할 수 없는 경우가 많습니다. 따라서 시스템 성능을 향상시키기 위한 주요 수단은 각 에이전트의 행동을 제어하는 프롬프트 (Prompt) 최적화가 됩니다.

하지만 MAS 의 프롬프트 최적화는 다음과 같은 세 가지 주요 난제에 직면해 있습니다:

1. 높은 평가 비용 (Expensive Evaluations): 후보 프롬프트 조합을 평가하려면 전체 MAS 워크플로우를 실행해야 하며, 이는 여러 LLM 호출을 필요로 하므로 평가 예산 (Budget) 이 매우 제한적입니다.
2. 토폴로지 유도 결합 (Topology-induced Coupling): 상류 (Upstream) 에이전트의 프롬프트 변경은 하류 (Downstream) 에이전트의 입력 분포를 변화시켜, 에이전트 간 의존성이 강하게 결합됩니다. 이로 인해 개별 에이전트를 독립적으로 최적화하는 것은 불안정합니다.
3. 조합적 폭발 (Combinatorial Explosion): 에이전트 수가 $N$이고 각 에이전트의 프롬프트 후보가 $|P_i|$개일 때, 전체 검색 공간은 $N$개의 곱집합으로 지수적으로 증가하여 완전 탐색이 불가능합니다.

기존의 단일 에이전트 프롬프트 최적화 방법 (OPRO, PromptBreeder 등) 은 에이전트 간 결합을 고려하지 않으며, 기존 다단계 최적화 방법 (MIPRO 등) 은 토폴로지 구조를 명시적으로 활용하지 않아 샘플 효율성이 낮습니다.

2. 제안 방법론: MASPOB

저자들은 MASPOB (Multi-Agent System Prompt Optimization via Bandits) 를 제안합니다. 이는 밴딧 (Bandit) 기반의 탐색 - 활용 (Exploration-Exploitation) 균형, 그래프 신경망 (GNN) 을 활용한 토폴로지 인식 모델링, 그리고 좌표 상승 (Coordinate Ascent) 을 통한 확장 가능한 검색을 통합한 프레임워크입니다.

핵심 구성 요소

1. 토폴로지 인식 대리 모델 (Topology-Aware Surrogate, GNN):

   • MAS 워크플로우를 방향성 비순환 그래프 (DAG) 로 모델링합니다.
   • 그래프 어텐션 네트워크 (GAT) 를 사용하여 각 에이전트의 프롬프트 임베딩을 노드 특징으로, 워크플로우 의존성을 인접 행렬로 활용합니다.
   • 메시지 전달 (Message Passing) 을 통해 에이전트 간의 상호작용과 프롬프트 변경이 시스템 전체에 미치는 영향을 학습하여, 전체 시스템 성능을 예측합니다. 이는 구조적 인덕티브 바이어스 (Structural Inductive Bias) 를 제공합니다.

2. 불확실성 기반 탐색 (Uncertainty-Driven Exploration, LinUCB):

   • 프롬프트 최적화를 컨텍스트 밴딧 (Contextual Bandit) 문제로 공식화합니다.
   • LinUCB (Linear Upper Confidence Bound) 알고리즘을 사용하여 평가 예산 내에서 최적의 조합을 찾습니다.
   • UCB 점수 = 예측 성능 ($\mu$) + 불확실성 보너스 ($\alpha \cdot \sigma$).
   • 불확실성 ($\sigma$) 은 정보 행렬 (Information Matrix) 을 통해 추정되며, 아직 충분히 평가되지 않은 영역 (Exploration) 을 탐색하도록 유도합니다.

3. 확장 가능한 조합 검색 (Scalable Combinatorial Search, Coordinate Ascent):

   • 전체 조합을 한 번에 최적화하는 대신, 좌표 상승 (Coordinate Ascent) 전략을 사용합니다.
   • 한 번에 하나의 에이전트 프롬프트만 변경하고 나머지는 고정하여 UCB 점수를 최대화하는 방식으로 반복합니다.
   • 이 접근법은 검색 복잡도를 지수적 ($O(\prod |P_i|)$) 에서 선형 ($O(\sum |P_i|)$) 으로 감소시켜, 실제 MAS 실행 비용이 아닌 GNN 추론 비용으로만 검색을 수행할 수 있게 합니다.

3. 주요 기여 (Key Contributions)

• 문제 공식화: 토폴로지 유도 결합과 이산적 조합 검색 공간을 가진 예산 제한 블랙박스 최적화 문제로 MAS 프롬프트 최적화를 공식화했습니다.
• MASPOB 프레임워크: GNN 대리 모델, LinUCB 기반 밴딧 탐색, 좌표 상승 전략을 결합하여 제한된 평가 예산 하에서 샘플 효율적인 최적화를 가능하게 했습니다.
• 성능 검증: 질문 답변, 코드 생성, 수학 추론 등 6 개의 벤치마크에서 기존 단일 에이전트 및 다중 에이전트 베이스라인 (AFlow, MIPRO 등) 을 일관되게 능가하는 결과를 달성했습니다.

4. 실험 결과 (Experimental Results)

• 벤치마크 성능: HotpotQA, DROP, HumanEval, MBPP, GSM8K, MATH 등 6 개 데이터셋에서 MASPOB 는 평균 **80.58%**의 정확도를 기록했습니다.
  • 기존 베이스라인 대비 평균 12.02% (IO 대비), 2.06% (AFlow 대비), 1.71% (MIPRO 대비) 의 성능 향상을 보였습니다.
  • 특히 복잡한 토폴로지를 가진 워크플로우에서도 AFlow 나 MIPRO 보다 우수한 일반화 능력을 입증했습니다.
• 수렴성: 평가 횟수가 증가함에 따라 검증 세트와 테스트 세트 모두에서 성능이 꾸준히 향상되었으며, 약 35 회 평가 시점에 최적 조합을 수렴하는 것을 확인했습니다.
• Ablation Study (성분 분석):
  • GNN 제거: GNN 을 MLP 로 대체하면 평균 성능이 2.31% 하락하여, 토폴로지 구조를 명시적으로 모델링하는 것이 에이전트 간 결합을 파악하는 데 필수적임을 입증했습니다.
  • 선형 vs 신경 불확실성: LinUCB(선형) 방식이 신경망 기반 불확실성 추정보다 더 빠른 수렴과 높은 성능을 보였습니다. 이는 데이터가 부족한 (Low-sample) 환경에서 선형 추정치가 더 안정적이기 때문입니다.
  • 좌표 상승 vs 전역 검색: 좌표 상승은 전역 검색과 유사한 성능을 내면서 실행 시간을 98~99.8% 단축했습니다.

5. 의의 및 결론 (Significance & Conclusion)

MASPOB 는 워크플로우 구조를 변경할 수 없는 실제 산업 환경 (High-stakes applications) 에서 MAS 의 성능을 극대화할 수 있는 실용적인 솔루션을 제공합니다.

• 구조 보존 최적화: 에이전트 간 의존성을 그래프 구조로 명시적으로 모델링함으로써, 단순한 프롬프트 튜닝을 넘어 시스템 전체의 조율 (Coordination) 을 최적화합니다.
• 비용 효율성: 고비용의 LLM 실행을 최소화하면서 GNN 과 밴딧 탐색을 통해 효율적으로 최적 조합을 찾습니다.
• 범용성: 다양한 LLM 백본 (GPT-4o-mini, Qwen3-32B 등) 과 다양한 작업 유형에서 일관된 성능 향상을 보여주어, 특정 모델에 종속되지 않는 일반적인 방법론임을 입증했습니다.

결론적으로 MASPOB 는 다중 에이전트 시스템의 프롬프트 최적화 분야에서 샘플 효율성과 토폴로지 인식이라는 두 가지 핵심 과제를 해결한 획기적인 접근법으로 평가됩니다.