PRISM: PRIvileged Sibling-guided Multi-Agent Topology Design

AI summary

Overview

Research area: LLM-based multi-agent systems (MAS), specifically the automated design of collaboration topologies, and post-training / on-policy distillation for graph-generating models.

Technical level: Advanced. The paper assumes familiarity with autoregressive sequence/graph generation, teacher forcing, KL-divergence objectives, on-policy distillation, and FiLM conditioning.

Scope (one sentence): PRISM is a post-training method that uses a second, training-only "sibling" graph that also solved a query to supervise an autoregressive multi-agent topology designer along trajectories the student itself samples, then discards that extra input at deployment.

What This Paper Is About

LLM multi-agent systems work better or worse depending on their collaboration topology — which agents participate and how they pass information — and recent methods learn a generator that builds such a topology from each task query, trained on graphs that were executed and verified to have solved that query. The problem is that this corpus records which graphs succeeded but never ranks them, so continued supervised fine-tuning has no per-query preference to learn from and drifts toward whatever topologies appear most often in the corpus, while teacher forcing also leaves the model unsupervised on partial graphs it reaches through its own deviations. PRISM's goal is to extract a per-query training signal from that same unranked corpus, without running any extra multi-agent executions and without changing the deployed model's architecture.

Key Contributions

1. A documented failure mode of continued SFT: the authors measure increased output concentration in an autoregressive topology designer under continued supervised fine-tuning, and relate it to accuracy changes across benchmarks (including a 12.60% drop relative to ARG-Designer on AQuA).

2. The PRISM method: successful sibling graphs are used as training-only privileged inputs for on-policy self-distillation, requiring no additional post-training LLM executions and no change to the deployed designer architecture.

3. Six-benchmark evaluation: on GSM8K, MultiArith, SVAMP, AQuA, MMLU and HumanEval, PRISM outperforms the strongest topology-generation baseline on each, under a single Qwen2.5-7B-Instruct execution stack.

4. Matched-teacher controls: shuffled, constant and alternative privileged channels isolate the contribution of query-matched graph information from generic on-policy correction or added parameters.

Main Findings

• Best on every benchmark: with Qwen2.5-7B-Instruct, PRISM scores 88.02 on GSM8K, 99.17 on MultiArith, 93.30 on SVAMP, 76.38 on AQuA, 82.32 on HumanEval and 73.94 on MMLU. The strongest baseline, ARG-Designer, scores 79.98, 98.83, 92.50, 75.20, 79.27 and 71.98 respectively.

• SFT is unstable across tasks: the SFT control beats or tracks ARG-Designer on five of six benchmarks (by at most 6.15% and never more than 0.10% below) but drops 12.60% on AQuA (62.60 versus 75.20).

• PRISM's margin over SFT varies by an order of magnitude: 0.72% on MMLU, 1.89% on GSM8K, and 13.78% on AQuA.

• SFT concentrates output, PRISM does not: top-1 topology share moves from 0.171 at the deployed starting checkpoint to 0.346 under SFT and 0.157 under PRISM. The number of distinct shapes is 22 at the cold start, 40 for SFT and 34 for PRISM. Resampling eight times, the query explains 1.1% of the cold start's output variety, 2.0% of SFT's and 5.8% of PRISM's.

• The modal crew does not explain the whole AQuA gap: SFT's most frequent role assignment is [Math Solver, Programming Expert, Programming Expert], emitted on 0.349 of GSM8K, 0.323 of MultiArith, 0.371 of SVAMP and 0.358 of AQuA questions. Its accuracy when emitted versus otherwise is +2.0% on GSM8K, +0.2% on MultiArith, -1.2% on SVAMP and -3.4% on AQuA.

• Difficulty-matched control favours PRISM in both groups: on AQuA, PRISM gains +18.7% (79.1 versus 60.4, McNemar p = 0.0001) where SFT emitted its default crew and +11.0% (p = 0.003) elsewhere. GSM8K gains +1.1% and +2.4%, SVAMP +3.0% and -0.3%, none significant at the 5% level at p = 0.68, 0.15, 0.16 and 0.92. On the 91 AQuA questions where SFT defaulted, PRISM emits 18 distinct topologies and repeats SFT's default crew on one question in eight.

• The privileged channel carries real information: the paired privilege gain reaches +2.45 nats for the real sibling channel while the shuffled control stays at zero; the FiLM norm rises to 3.10 with sibling inputs and 1.15 with shuffled inputs, so the zero gain is a genuine null rather than an untrained branch.

• Content beats procedure: a constant-input teacher beats the shuffled teacher on all six benchmarks (83.86 versus 82.58 average), and PRISM beats the constant teacher on all six by +1.66% on average (sign test p = 0.031). On AQuA, SFT is 62.60, the constant teacher 71.65 (+9.05% from on-policy distillation) and PRISM 76.38 (+4.73% from content). The constant teacher improves on SFT on only three benchmarks, losing 1.11% on MMLU.

• Sibling graph is the best channel: it wins on all six benchmarks, with margins over the best competitor of 3.26% on GSM8K and 14.96% on AQuA. A "context" channel (question plus solution) yields approximately 0 nats, moves the embedding 5.8 times less than shuffled context (node-level KL 0.0101 versus 0.0584), and distilling from it costs 5.84% on GSM8K and 14.96% on AQuA.

• Reverse KL is the best divergence on average: 85.52 versus 84.85 for forward KL and 83.90 for JSD (β = 0.5); forward KL is slightly better on HumanEval (+0.61%) and MMLU (+0.33%), while reverse KL leads forward KL by 1.89% and 2.10% on GSM8K and SVAMP and JSD by 2.80% and 3.00%.

• Near-identical deployment cost: PRISM averages 3.04 agents and 4.04 LLM calls per query at K = 1, against 3.00 agents and 4.00 calls for both ARG-Designer and SFT — 2.56% and 3.24% more accuracy for 1.0% more calls.

• Corpus is uninformative at the margin but not conditionally: across the four MixChain families with recorded attempts, solve rates span 58.5% to 59.5%, a spread of 1%, despite differing in edge density by a factor of two; corpus shares instead track how often each family was sampled (chain 35.1%, star 29.4%, layered 18.4%, fully-connected 9.7%, mesh 7.4%).

Methodology in Plain English

The starting point. The authors build on ARG-Designer, an autoregressive generator that constructs a collaboration graph for a query one decision at a time: it picks a role for each new agent (with a learnable END token that stops construction) and then decides which earlier agents that agent receives input from. Supervision comes from a "verified corpus": for each training query, several blueprint configurations (chain, star, layered, fully-connected, mesh) were executed, and only the graphs that actually solved the query were kept. For MixChain this yields 2,947 verified graphs over 1,014 distinct queries and 25 distinct topologies, from 5,836 successful executions out of 9,902 attempts.

The diagnosis. That corpus says a graph worked, never that it worked better than another graph for the same query. Continuing to fine-tune on it therefore pushes the model toward frequent shapes rather than query-appropriate ones, and because training only ever shows recorded construction paths, the model is never corrected on partial graphs it reaches by its own choices.

Stage 1 — build a teacher with a private hint. The authors copy the trained generator to make a teacher with no shared parameters, and add a small branch that reads a privileged input: another graph that also solved the same query, converted to text and embedded with the same frozen all-MiniLM-L6-v2 encoder. That embedding modulates the teacher's conditioning vector through a FiLM transform with rank r = 16, with the modulation output maps zero-initialised so the teacher starts out behaving exactly like the student. The branch adds about 20k parameters versus about 592k for a full-width channel at d = 384. A sibling is drawn uniformly among the other distinct graphs that solved that query, re-drawn every epoch; queries with no distinct sibling get a zero vector rather than another query's graph. The teacher trains for 40 epochs at learning rate 10⁻⁴.

Stage 2 — correct the student where it actually goes. The student samples M = 4 full topologies per query at temperature τ = 1.0, and both models re-score those exact trajectories under teacher forcing — the student without the hint, the teacher with it. Because both predict over the same small categorical decision spaces (role choices plus END, and binary edge decisions), the per-decision divergence is computed exactly rather than estimated. The loss is a KL term (reverse KL by default, with α = 0.2 weighting nodes versus edges) mixed with a supervised anchoring loss on a verified graph, weighted λ = 0.3. Stage 2 runs 30 epochs at 10⁻⁴ with 16 queries per step. Sampled graphs are scored as sampled, never repaired first; repairs for executability are applied only at evaluation. Gradients flow only through re-scoring, and the teacher is frozen and then thrown away.

Stage 3 — deploy the plain model. The deployed student is architecturally identical to ARG-Designer and sees only the query.

Evaluation design. All arms are controlled re-runs on Qwen2.5-7B-Instruct with K = 1 round plus a summariser, using the same pipeline. Training and validation are split by query with a fixed seed and 20% held out. The corpus for the four mathematical benchmarks comes from MixChain, MMLU from its own questions, and code generation from MBPP; HumanEval is evaluation-only. No hyperparameters are tuned per benchmark.

Why This Matters

Impact on research. The paper reframes post-training of topology generators as a distillation problem rather than a likelihood-fitting problem, and shows that a corpus of unranked successes can still supply a per-query signal if it is injected as privileged information during training only. It also gives a reusable diagnostic toolkit — top-1 topology share, query-variance against a permutation null, privilege gain in nats, and shuffled/constant teacher controls — for separating "the model learned something query-specific" from "the model learned a better constant" or "the extra parameters did the work."

Potential real-world applications:

• Automatically composing solver/reviewer agent crews for math and quantitative word problems, where PRISM shows its largest reported gain (AQuA, 76.38).
• Code-generation assistants that decide per request whether to route through a single agent or a multi-agent pipeline (HumanEval, 82.32).
• General-knowledge and exam-style question answering with cost-aware agent orchestration (MMLU, 73.94, at 4.04 LLM calls per query).
• Any setting where a log of previously successful configurations exists but is unranked — incident response runbooks, data-processing pipelines, retrieval workflows — since the method needs no new executions to post-train.

Industry relevance. The deployment cost is essentially unchanged (3.04 versus 3.00 agents per query), the post-training itself adds no agent executions because rollouts sample graph structure rather than running agents, and the privileged branch is removed at inference, so serving infrastructure does not change. That combination — no extra inference cost, no extra execution cost, no architecture change — is what makes the technique plausible for production systems that already ship a fixed topology generator.

Future Directions

• Test the crew-assignment effect directly. The authors state their results do not establish that query-specific crew assignment explains the gains; they propose evaluating that assignment effect on representative held-out questions.
• Compare post-training methods at matched compute. PRISM's Stage 1 (40 epochs) plus Stage 2 (30 epochs) does not match the SFT control's total training computation; the paper explicitly calls for matched-compute comparisons.
• Test transfer across backbones. All reported results use Qwen2.5-7B-Instruct; whether the sibling-guided scheme holds for other models is stated as future work.
• Establish whether a ranking among successful graphs is recovered. The paper is explicit that it does not show PRISM recovers a preference ordering over the graphs that solved a query, only that sibling content contributes beyond the distillation procedure.

Target Audience

Researchers and advanced practitioners working on LLM multi-agent systems, agent orchestration and automatic workflow/pipeline design; machine-learning engineers interested in on-policy distillation, privileged-information training and teacher-student methods; and applied scientists who need to post-train a structured generator from a corpus of unranked successes without adding inference cost. Readers should be comfortable with autoregressive generation, KL divergence variants and standard distillation terminology.

Authors’ abstract

Multi-agent systems (MAS) based on large language models (LLMs) have become a powerful way to address complex problems, and their effectiveness depends critically on the collaboration topology that decides which agents participate and how they communicate. Recent work learns to generate this topology from the task query, fitting the generator to corpora of execution-verified graphs. The post-training method for such generators remains a limitation: a corpus lists the graphs that solved a query but never ranks them, leaving continued supervised fine-tuning (SFT) with no per-query preference to exploit. Without such a preference, the generator drifts toward the topologies the corpus contains most often rather than those that suit each query, so a gain on one task can come at a cost on another. We propose PRISM, which recovers per-query signal from the same corpus by learning with privileged information: a teacher reads, through a zero-initialized low-rank branch, a sibling graph that also solved the query and is unavailable at deployment, and the student matches its per-decision targets on trajectories the student samples itself. Across six diverse benchmarks PRISM improves over every topology-design baseline evaluated, and matched-teacher controls support the contribution of privileged query information. Code is available in this link.