Source-linked AI summary
Codebook Agent: Amortized Topology Design for LLM Multi-Agent Systems
Jinxi Yu, Yubei Li, Eric Hanchen Jiang, Zhi Zhang, Dong Liu, Wenxiao Zhao, Levina Li, Kai-Wei Chang, Ying Nian Wu
TL;DR
Existing topology designers search the full adjacency space and often score candidates with edge-count-based graph networks, despite evidence that useful topologies form a short list and structural cost can misalign with measured tokens. Codebook Agent discretizes successful topologies, predicts query-conditioned codes, and reranks decoded candidates with an execution-grounded MLP proxy. Across six benchmarks, it leads prior methods while reducing design latency and token use.
Problem
Current query-conditioned topology designers treat communication design as adjacency-space generation followed by structural-cost graph scoring, although useful topologies may be limited and measured token cost can differ from edge count.
Method
Codebook Agent uses a query-independent 16-entry codebook of successful topologies, a reward-weighted MLP code predictor, and a flattened-adjacency MLP proxy trained on utility and normalized token cost.
Results
Codebook Agent leads all six benchmarks, averages 84.62 versus 83.02 for GTD, generates topologies in 2.4 ms, and reduces LLM tokens by 21.9–33.2%.
Takeaways & Limitations
Topology design can be amortized as fast selection from a small discrete set scored by measured inference cost rather than iterative graph generation.
Takeaways & Limitations
The benchmark setup uses homogeneous profiles in six of seven design-axis settings, with HumanEval as the heterogeneous four-role exception.
Abstract
from arXiv · showhide
Adapting the communication topology of an LLM multi-agent system to each query improves both accuracy and efficiency, yet current designers treat this as conditional graph generation: a variational, autoregressive, or diffusion decoder searches the $N \times N$ adjacency space, and a graph-network proxy trained on utility and a structural cost such as edge count ranks the sampled candidates. We argue that this formulation is misaligned with the problem. Empirically, topologies that survive a reward filter collapse to about six distinct graphs even when the codebook capacity grows from 8 to 64; edge count is negatively correlated with measured token consumption (Pearson $r \approx -0.4$), so sparsifying the graph makes inference more expensive; and a message-passing scorer over agent-profile nodes is adjacency-invariant whenever agents share a profile---the default configuration of published benchmarks---so it cannot rank candidates at all in that regime. These three facts motivate Codebook Agent: a vector-quantized autoencoder compresses successful topologies into a query-independent 16-entry codebook; a reward-weighted MLP maps the query embedding to a distribution over codes; and an MLP proxy that reads the flattened adjacency, regressed on measured utility and per-task normalized token cost, reranks the top decoded candidates in a single batched forward pass. With no iterative search and no message passing at test time, Codebook Agent is the most accurate method on all six benchmarks we compare (84.6 average against 83.0 for the strongest prior designer), emits a topology in 2.4 ms, and uses 21.9--33.2% fewer LLM tokens.
1 Introduction
LLM multi-agent topology design is miscast as large-scale graph generation: useful topologies collapse to a short list, edge count misrepresents token cost, and message-passing scorers can become topology-blind. Codebook Agent instead selects among discretized successful topologies with feed-forward predictors and a measured-cost proxy.
- Motivation: Communication topology affects both multi-agent accuracy and token consumption, motivating learned per-query topology design.The topology determines which agents see which messages.
- Motivation: About six distinct topologies survive reward filtering as codebook capacity grows from 8 to 64, while the best fixed topology remains within 1.4 accuracy points of generated ones.This suggests adjacency-generation capacity is not where the useful design space primarily resides.
- Motivation: Edge count correlates negatively with measured token consumption at r ≈−0.4, because sparse communication yields longer completions.Thus, minimizing edge count can favor more expensive inference.
- Motivation: Message-passing scorers are adjacency-invariant on homogeneous teams, assigning every candidate the same score and eliminating candidate ranking.Homogeneous profiles dominate published benchmark configurations.
- Approach: Codebook Agent indexes successful topologies into 16 codes, predicts codes from query embeddings, and reranks decoded candidates with an MLP proxy using measured utility and normalized token cost.The test-time path uses one batched proxy call without iterative search or message passing.
- Results: Codebook Agent is most accurate on all six benchmarks, generates topologies in 2.4 ms, and reduces token consumption by 21.9–33.2%.The comparison covers single-agent prompting, multi-agent collaboration, and learned topology designers.
2 Related Work
Prior work progresses from fixed or hand-designed communication structures to optimization, pruning, and query-conditioned topology generators. Codebook Agent shares per-query selection with the latest generators but replaces adjacency generation and graph-network scoring with selection from a short list and measured-cost scoring.
- Communication topologies for LLM agents: Earlier multi-agent frameworks fixed communication topologies by hand, while later systems optimized edges, selected agents, searched designs, pruned structures, or trained agents for efficiency.These approaches span policy-gradient learning, offline search, sparsification, and efficiency-oriented training.
- Communication topologies for LLM agents: Query-conditioned generators are the closest prior work because they emit one topology per query while modeling adjacency generation and graph-network scoring.Codebook Agent retains the problem setting but changes the representation and scorer.
- Graph generation and discrete latents: Iterative graph generators, including autoregressive and score-based models, inherit sampling cost even when faster samplers reduce it.The paper uses vector quantization because it needs one small graph per query under a latency constraint.
- Cost of multi-agent inference: Multi-agent inference can become less efficient as calls grow, motivating measured token count rather than a structural surrogate as the topology-design cost objective.Prior work reports diminishing or inverted returns and coordination-related failures.
3 Problem Setup and Background
The paper formalizes topology selection as query-conditioned optimization over binary directed adjacency matrices, then separates candidate generation from candidate scoring. Its setup uses execution records and contrasts iterative graph-network incumbents with the Codebook Agent pipeline.
- Topology design: A topology is a directed binary adjacency matrix without self-loops, and executing it on query q returns utility u(A, q) and token cost τ(A, q).The objective is R = u −λ˜τ with λ = 0.1 and normalized token cost.
- Topology design: The training records contain 300 tuples per benchmark from 50 tasks executed under six fixed topologies with real LLM agents.Each tuple includes adjacency, query embedding, utility, and token cost.
- Two design axes: Candidate generation and scoring are independent design axes: generators map query embeddings to topologies, while scorers rank candidates.Incumbent generators use variational, autoregressive, or diffusion-style adjacency decoding.
- Codebook Agent: Codebook Agent instead compresses successful topologies into a query-independent codebook, predicts codes, and scores decoded candidates with a one-pass proxy.Its proxy uses measured utility and per-task normalized token cost, with an auxiliary structure-token head used only during training.
- Two design axes: Incumbent diffusion-style generation takes 301 to 396 ms per query and uses a message-passing scorer over profile-embedded nodes with utility and structural-cost heads.The structural cost is most commonly edge count |E|.
4 Method
Codebook Agent amortizes topology design into a query-conditioned codebook and execution-grounded proxy, replacing iterative graph search with feed-forward candidate generation and reranking. The method uses measured, normalized token cost and flattened adjacency inputs to select among successful topologies.
- 4.2 Reward-Weighted Code Prediction: A reward-weighted MLP maps each query condition to a distribution over codes, concentrating probability on codes whose topologies earned high composite reward.The composite reward uses measured utility and token count normalized by the mean token count for the same task, with λ = 0.1.
- 4.1 Topology Codebook: Codebook Agent represents successful topologies with a vector-quantized autoencoder and a 16-entry codebook.The topology is encoded from flattened off-diagonal adjacency entries and decoded as a directed binary graph with a zero diagonal.
- 4.1 Topology Codebook: The codebook is trained only on records with utility uj > 0.5, so its codes store topologies that solved their tasks.The encoder and decoder are query-independent; query dependence is introduced by the predictor.
- 4.3 Execution-Grounded MLP Proxy: The execution-grounded proxy reads flattened adjacency and regresses measured utility and normalized token cost rather than edge count.Flattened adjacency is necessary because message passing over homogeneous agent profiles is independent of topology, while edge count correlates negatively with measured tokens at r ≈−0.4.
- 4.4 Test-Time Selection: At test time, the system decodes the top M = 5 codes, removes duplicate topologies, and scores survivors in one batched forward pass.Generation uses one predictor pass, at most M decoder passes, and one dense proxy pass, with no sampling loop or test-time message passing.
5 Experiments
Experiments evaluate whether amortized codebook selection preserves accuracy while reducing topology-design and inference costs across benchmarks, ablations, and a backbone transfer. Codebook Agent leads prior designers while using one-pass selection and fewer tokens.
- Main results: Codebook Agent achieves the best accuracy on every Table 1 benchmark, averaging 84.62 versus 83.02 for prior designer GTD.It exceeds GTD by 1.3–3.1 accuracy points across GSM8K, MATH, MultiArith, SVAMP, MBPP, and HumanEval.
- Main results: The best fixed topology remains within 1.4 accuracy points of every generated topology across seven settings, at comparable token usage.On MATH, fully connected is above every generated configuration.
- Main results: 21.9–33.2% fewer tokens and 2.4–2.5 ms generation establish Codebook Agent’s accuracy–efficiency advantage over the full incumbent pipeline.Iterative generators require 301–396 ms, while Codebook Agent uses one predictor pass, limited decodes, and one batched MLP call.
- Ablations: MLP reranking uses 927/546 tokens on GSM8K/HumanEval, compared with 1249/750 for random and 1711/918 for GNN reranking.Accuracy stays within 1.5 points, indicating that the selector primarily determines token cost in this ablation.
- Ablations: For every K≥8, the encoder uses at most six codes, while increasing K changes accuracy by at most 1.5 points on GSM8K and 2.5 on HumanEval.The ablation uses K ∈ {4, . . . , 64} under a fixed training budget.
- Backbone transfer: With Qwen-3-8B, Codebook Agent still leads the three-benchmark transfer at 74.0 versus 72.7 for GTD and 72.1 for G-Designer.The transfer evaluates GSM8K, MATH, and MMLU under the same agent teams and protocol.
6 Conclusion
Codebook Agent finds that effective topologies form a short list and selects among them with a lightweight proxy, leading across benchmarks while reducing design latency and token use.
- Codebook Agent identifies effective multi-agent topologies as a short list rather than a large design space.
- Denser graphs need not reduce token consumption, challenging edge count as a structural cost surrogate.
- Message-passing critics can miss topology structure when teams are homogeneous.
- Codebook Agent leads all six benchmarks and Qwen-3-8B transfer while cutting design latency to 2.4 ms and LLM tokens by 21.9–33.2%.