Official code for the ICML 2025 paper “KABB: Knowledge-Aware Bayesian Bandits for Dynamic Expert Coordination in Multi-Agent Systems”.

[ English | 中文 ]

KABB (Knowledge-Aware Bayesian Bandits) is a dynamic expert coordination framework for multi-agent systems, featuring:

1. Knowledge Distance Model — Semantic matching between experts and tasks.
2. Dual Adaptation Mechanism — Continuous optimization of expert representation and selection.
3. Knowledge-Aware Thompson Sampling — Efficient expert selection in Bayesian MAB with knowledge distance.

See the paper for theoretical details and full experiments.

1. Clone the repo:

   git clone https://github.com/your_org/KABB.git cd KABB

2. (Optional) Create a virtual environment:

   python3 -m venv .venv && source .venv/bin/activate

3. Install dependencies:

   pip install -r requirements.txt

4. Configure environment variables:

   cp .env.example .env # Edit .env and fill in your API keys

Run a sample math task with KABB:

python scripts/run_kabb.py \ --config configs/config_math_template.yaml \ --question "What is the value of (7/8)^3 * (7/8)^-3?"

configs/config_math_template.yaml provides a minimal working config, including:

• system_prompts: System prompts for different scenarios
• domain_inference_settings: Domain priors, samples, and key symbols
• experts_pool: List of experts (model, temperature, max tokens, etc.)
• llm: API key placeholder (use environment variable)

To extend to other domains:

1. Add a new domain entry with prior and typical_samples
2. Prepare a set of expert models for the domain
3. Specify the new config in the script

Reproduce main results with:

python scripts/run_kabb.py --config configs/config_math_template.yaml --question "..."

For full benchmarks, see the run_*.py scripts and refer to the paper appendix.

If you use this project, please cite:

@misc{zhang2025kabbknowledgeawarebayesianbandits, title={KABB: Knowledge-Aware Bayesian Bandits for Dynamic Expert Coordination in Multi-Agent Systems}, author={Jusheng Zhang and Zimeng Huang and Yijia Fan and Ningyuan Liu and Mingyan Li and Zhuojie Yang and Jiawei Yao and Jian Wang and Keze Wang}, year={2025}, eprint={2502.07350}, archivePrefix={arXiv}, primaryClass={cs.AI}, url={https://arxiv.org/abs/2502.07350}, }

• Fork and create a new branch
• Link related issues in PRs

MIT. See LICENSE.