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

TL;DR

MASPOB (Multi-Agent System Prompt Optimization via Bandits), a novel sample-efficient framework based on bandits that achieves state-of-the-art performance, consistently outperforming existing baselines.

Abstract

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 (5)

• 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.
• 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.
• Figure 4: Performance and runtime comparison between coordinate ascent and global search. The figure illustrates optimization trajectories and time costs on selected benchmarks.
• Figure 5: Uncertainty convergence of LinUCB-style linear uncertainty and neural uncertainty over optimization rounds. With only 45 exploration rounds, the linear uncertainty decreases by 71.68%, whereas the neural uncertainty decreases by 22.48%, suggesting that neural uncertainty may require more exploration rounds to reach a comparable level of confidence.