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

Zhi Hong, Qian Zhang, Jiahang Sun, Zhiwei Shang, Mingze Kong, Xiangyi Wang, Yao Shu, Zhongxiang Dai
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:43744-43774, 2026.

Abstract

Large Language Models (LLMs) have achieved substantial success in real-world applications, particularly as the cognitive backbone of Multi-Agent Systems (MAS) for orchestrating complex workflows. Since many deployments preclude workflow modifications while MAS performance is highly prompt-sensitive, prompt optimization becomes a critical strategy for improvement. 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. Our code is available at https://github.com/HZ1008/MASPOB.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-hong26l, title = {{MASPOB}: Bandit-Based Prompt Optimization for Multi-Agent Systems with Graph Neural Networks}, author = {Hong, Zhi and Zhang, Qian and Sun, Jiahang and Shang, Zhiwei and Kong, Mingze and Wang, Xiangyi and Shu, Yao and Dai, Zhongxiang}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {43744--43774}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/hong26l/hong26l.pdf}, url = {https://proceedings.mlr.press/v306/hong26l.html}, abstract = {Large Language Models (LLMs) have achieved substantial success in real-world applications, particularly as the cognitive backbone of Multi-Agent Systems (MAS) for orchestrating complex workflows. Since many deployments preclude workflow modifications while MAS performance is highly prompt-sensitive, prompt optimization becomes a critical strategy for improvement. 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. Our code is available at https://github.com/HZ1008/MASPOB.} }
Endnote
%0 Conference Paper %T MASPOB: Bandit-Based Prompt Optimization for Multi-Agent Systems with Graph Neural Networks %A Zhi Hong %A Qian Zhang %A Jiahang Sun %A Zhiwei Shang %A Mingze Kong %A Xiangyi Wang %A Yao Shu %A Zhongxiang Dai %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-hong26l %I PMLR %P 43744--43774 %U https://proceedings.mlr.press/v306/hong26l.html %V 306 %X Large Language Models (LLMs) have achieved substantial success in real-world applications, particularly as the cognitive backbone of Multi-Agent Systems (MAS) for orchestrating complex workflows. Since many deployments preclude workflow modifications while MAS performance is highly prompt-sensitive, prompt optimization becomes a critical strategy for improvement. 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. Our code is available at https://github.com/HZ1008/MASPOB.
APA
Hong, Z., Zhang, Q., Sun, J., Shang, Z., Kong, M., Wang, X., Shu, Y. & Dai, Z.. (2026). MASPOB: Bandit-Based Prompt Optimization for Multi-Agent Systems with Graph Neural Networks. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:43744-43774 Available from https://proceedings.mlr.press/v306/hong26l.html.

Related Material