MASPOB: Bandit-Based Prompt Optimization for Multi-Agent Systems with Graph Neural Networks
Authors: Zhi Hong, Qian Zhang, Jiahang Sun, Zhiwei Shang, Mingze Kong, Xiangyi Wang, Yao Shu, Zhongxiang Dai
Organizations: The Chinese University of Hong Kong, Shenzhen · South China University of Technology · Ritsumeikan University · The Hong Kong University of Science and Technology (Guangzhou)
Large Language Models (LLMs) have achieved great success in many real-world applications, especially the one serving as the cognitive backbone of Multi-Agent Systems (MAS) to orchestrate complex workflows in practice. Since many deployment scenarios preclude MAS workflow modifications and its performance is highly sensitive to the input prompts, prompt optimization emerges as a more natural approach to improve its performance. However, real-world prompt optimization for MAS is impeded by three key challenges: (1) the need of sample efficiency due to prohibitive evaluation costs, (2) topology-induced coupling among prompts, and (3) the combinatorial explosion of the search space. To address these challenges, we introduce MASPOB (Multi-Agent System Prompt Optimization via Bandits), a novel sample-efficient framework based on bandits. By leveraging Upper Confidence Bound (UCB) to quantify uncertainty, the bandit framework balances exploration and exploitation, maximizing gains within a strictly limited budget. To handle topology-induced coupling, MASPOB integrates Graph Neural Networks (GNNs) to capture structural priors, learning topology-aware representations of prompt semantics. Furthermore, it employs coordinate ascent to decompose the optimization into univariate sub-problems, reducing search complexity from exponential to linear. Extensive experiments across diverse benchmarks demonstrate that MASPOB achieves state-of-the-art performance, consistently outperforming existing baselines.
Figures & tables
Figure 1: The MASPOB framework. (1) Initialization: Constructs agent topology and generates prompt embeddings. (2) Search: Selects optimal prompts via Coordinate Ascent, balancing exploitation (GNN prediction) and exploration (Linear UCB uncertainty). (3) Update: Refines the GNN model and information matrix using execution feedback.
Figure 2: Performance evaluation of prompt optimization methods. (a) Detailed comparison across six diverse benchmarks including question answering (HotpotQA, DROP), code generation (HumanEval, MBPP), and mathematical reasoning (GSM8K, MATH). (b) Overall ranking based on average performance. MASPOB demonstrates superior performance with an average improvement of 12.02% over the IO baseline.
QA
Code
Math
Method
HotpotQA
DROP
HumanEval
MBPP
GSM8K
MATH
Avg.
IO (GPT-4o-mini)
60.36 ±0.48
53.09 ±0.36
89.31 ±2.02
69.11 ±1.19
87.80 ±0.44
51.71 ±0.43
68.56
CoT ( Wei et al., 2022 )
67.62 ±0.48
58.27 ±0.42
89.57 ±1.16
69.89 ±1.19
88.34 ±0.18
52.47 ±1.25
71.03
ReAct ( Yao et al., 2022 )
65.61 ±0.22
67.25 ±0.85
87.79 ±1.32
66.08 ±1.48
88.91 ±0.24
52.61 ±0.62
71.38
PromptBreeder ( Fernando et al., 2023 )
68.76 ±0.18
71.85 ±0.87
88.80 ±2.45
70.38 ±0.36
91.97 ±0.80
52.13 ±0.97
73.98
Instinct ( Lin et al., 2023 )
69.92 ±0.17
71.90 ±0.53
90.08 ±2.76
70.23 ±1.14
92.64 ±0.39
52.40 ±1.52
74.53
Table 1: Performance comparison across six benchmarks. We report mean accuracy (%) ± standard deviation over three runs. Bold indicates the best result for each benchmark.
Method
HotpotQA
DROP
HumanEval
Avg.
Aflow
73.38 ± 0.36
80.28 ± 1.07
89.57 ± 0.44
81.08
MIPRO
74.03 ± 0.38
79.99 ± 0.78
86.77 ± 1.16
80.26
MASPOB
74.43 ± 0.37
81.55 ± 1.29
90.08 ± 0.76
82.02
Table 2: Performance on more complex MAS structures. We generate structures with greater topological complexity using AFlow (8, 7, and 7 agents for HotpotQA, DROP, and HumanEval, respectively), compared to the original simpler topologies (3, 2, and 3 agents). All other experimental settings are kept unchanged. We report mean accuracy (%) ± standard deviation over three runs. Bold indicates the best result for each benchmark.
Figure 3: Optimization convergence on validation and test sets. The curves show the average validation accuracy, computed by averaging over every five rounds, and the test accuracy at rounds 5, 10, …, 50. For each selected test combination, the accuracy at these rounds is evaluated and averaged over three runs.
Method
HotpotQA
DROP
HumanEval
MBPP
GSM8K
MATH
Avg.
AFlow
73.42 ±0.38
79.48 ±0.12
91.09 ±0.44
79.96 ±0.67
93.36 ±0.43
53.83 ±0.40
78.52
MIPRO
74.37 ±1.07
79.13 ±0.59
91.35 ±0.44
80.65 ±0.36
92.80 ±0.42
54.90 ±0.61
78.87
MASPOB
75.43 ±0.27
82.28 ±0.55
94.15 ±0.44
80.65 ±0.29
93.90 ±0.15
57.05 ±0.51
80.58
w/o GNN
74.38 ±0.44
79.20 ±0.64
90.33 ±1.08
79.18 ±0.77
92.95 ±0.31
53.57 ±0.83
78.27
Δ (GNN Gain)
+1.05
+3.08
+3.82
+1.47
+0.95
+3.48
+2.31
Table 3: Ablation study replacing the GNN with an MLP. We replace the GNN module with a multi-layer perceptron (MLP) while keeping all other components unchanged. We report mean accuracy (%) ± standard deviation over three runs. Bold indicates the best result for each benchmark.
Method
DROP
HumanEval
MATH
Avg.
Aflow
84.73 ± 0.48
92.11 ± 0.44
69.97 ± 0.32
82.27
MIPRO
85.28 ± 0.02
93.64 ± 0.54
70.70 ± 1.25
83.21
Our
85.31 ± 0.47
94.15 ± 0.44
70.84 ± 0.48
83.43
Table 4: Generalization to Qwen-3-32B. Results using Qwen-3-32B as the backbone LLM, replacing GPT-4o-mini while keeping all other settings unchanged. We report mean accuracy (%) ± standard deviation over three runs. Bold indicates the best result for each benchmark.
Prompt Domain
DROP
MATH
Avg.
MIPRO domain
82.81 ± 0.15
57.45 ± 0.61
70.13
MASPOB domain
82.28 ± 0.55
57.05 ± 0.51
69.67
Table 5: Prompt-domain robustness. We replace the MASPOB prompt domain with MIPRO’s while keeping the MASPOB optimizer unchanged. We report mean accuracy (standard deviation over three runs. Bold indicates the best result for each benchmark.
Method
HotpotQA
DROP
Perf.(%)
Time(s)
Perf.(%)
Time(s)
Coordinate
75.43 ± 0.27
15.9
82.28 ± 0.55
7.7
Global
75.72 ± 0.29
8801
82.76 ± 0.27
392
Δ
-0.29
99.8% ↓
-0.48
98.0% ↓
Table 6: Performance and runtime comparison between Coordinate Ascent and Global Search. We replace the coordinate ascent strategy with an exhaustive global search on the HotpotQA and DROP benchmarks, while keeping all other experimental settings unchanged. We report mean accuracy (%) ± standard deviation over three runs.
Figure 4: Performance and runtime comparison between coordinate ascent and global search. The figure illustrates optimization trajectories and time costs on selected benchmarks.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Dataset
Validation
Test
DROP
100
800
HotpotQA
100
800
GSM8K
150
1,055
MATH
100
486
MBPP
86
341
HumanEval
33
131
Appendix
Table 7: Dataset splits used in our experiments. Validation samples are used for prompt optimization, while test samples are held out for final evaluation.
Category
Parameter
Value
Optimization
Maximum optimization rounds
45
Prompt variants per agent
20
Pretraining rounds
5
Search strategy
Coordinate
Random seed
42
UCB
Exploration coefficient α
0.2
Appendix
Table 8: Hyperparameter configurations for MASPOB.
Table 13
Figure 5: Uncertainty convergence over optimization rounds. Linear uncertainty decreases by 71.68%, whereas neural uncertainty decreases by 22.48% within 45 rounds, suggesting that neural uncertainty may require more exploration to reach comparable confidence.
Large language model (LLM)-based Multi-agent systems (MAS) have shown promise in tackling complex collaborative tasks, where agents are typically orchestrated via role-specific prompts. While the quality of these prompts is pivotal, jointly optimizing them across interacting agents remains a non-trivial challenge, primarily due to the misalignment between local agent objectives and holistic system goals. To address this, we introduce MASPO, a novel framework designed to automatically and iteratively refine prompts across the entire system. A core innovation of MASPO is its joint evaluation mechanism, which assesses prompts not merely by their local validity, but by their capacity to facilitate downstream success for successor agents. This effectively bridges the gap between local interactions and global outcomes without relying on ground-truth labels. Furthermore, MASPO employs a data-driven evolutionary beam search to efficiently navigate the high-dimensional prompt space. Extensive empirical evaluations across 6 diverse tasks demonstrate that MASPO consistently outperforms state-of-the-art prompt optimization methods, achieving an average accuracy improvement of 2.9. We release our code at https://github.com/wangzx1219/MASPO.
Zhexuan Wang, Xuebo Liu, Li Wang +4
Institute of Computing and Intelligence, Harbin Institute of Technology, Shenzhen, China.
Multi-agent systems (MAS) offer a scalable path forward for agentic AI, comprising multiple LLM-based agents, each assigned a system prompt and a position within a workflow that governs inter-agent coordination and output aggregation. System prompts thus form a critical and accessible optimization surface: they specify agents' roles and behaviors, enabling system-level improvements without model finetuning. Although prompt optimization has shown substantial potential for single LLMs, extending it to MAS poses distinct challenges, notably an exponentially growing search space. It remains unclear whether, when, and by how much prompt optimization improves MAS performance, and how sensitive such gains are to system configuration. In this work, we systematically study system-prompt optimization across a broad range of MAS setups varying in task, workflow, communication protocol, and team size, benchmarking two prompt optimizers that naturally extend state-of-the-art single-agent methods. The results reveal its potential to unlock significant gains while exposing open challenges, characterizing when and how much prompt optimization helps across diverse MAS settings.
Juyang Bai, Laixi Shi
Department of Electrical and Computer Engineering, Johns Hopkins University, Baltimore, MD 21218, USA.
Large language model (LLM)-based multi-agent systems (MAS) rely on role-specific prompts and execution models to solve complex tasks collaboratively. However, prompt optimization typically keeps model assignments fixed, leaving differences in role requirements and their associated inference costs outside the search. We present UniMASO, a Unified framework for heterogeneous Multi-Agent System Optimization through the joint adaptation of prompts and model assignments. UniMASO evaluates candidate updates through local, downstream, and global feedback, using cost-aware Pareto search to retain useful quality-cost trade-offs. Topology-aware co-evolution then coordinates these updates according to agent roles and dependencies. Cross-level trajectory memory further supplies historical successes and failures to guide subsequent optimization. Experiments on five benchmarks and five execution topologies show that UniMASO achieves competitive task performance at substantially lower deployment cost than the strongest homogeneous configurations. Under centralized execution, it improves the average score from 69.79 to 70.40 while reducing token-weighted deployment cost by 85.2% relative to DeepSeek-V4-Pro with MASPO. These results support joint prompt-model adaptation as an effective approach to efficient multi-agent collaboration. The code is available at https://github.com/Achilles008/UniMASO_main.
Mingxiu Cai, Haoyang Xian, Xinyu Ma +4
Institute of Computing and Intelligence, Harbin Institute of Technology, Shenzhen