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.