Source-linked AI summary
OFA-MAS: One-for-All Multi-Agent System Topology Design based on Mixture-of-Experts Graph Generative Models
Shiyuan Li, Yixin Liu, Yu Zheng, Mei Li, Quoc Viet Hung Nguyen, Shirui Pan
TL;DR
MAS performance depends on collaboration topology, while existing topology-design approaches often rely on domain-specific models. OFA-MAS generates task-adaptive collaboration graphs with a unified model and achieves a 93.02% average score, outperforming all baselines.
Problem
MAS effectiveness depends on collaboration topology, while existing design models often lack mechanisms to incorporate task semantics into role selection across domains.
Source: §1 · pp. 1–2Method
Results
93.02% average score; OFA-MAS outperforms all baselines across the evaluated benchmarks.
Source: §4.2 · p. 7Takeaways & Limitations
No single fixed MAS topology excels across all tasks, so topology performance varies with task demands.
Source: §4.2 · p. 7
Abstract
from arXiv · showhide
Multi-Agent Systems (MAS) offer a powerful paradigm for solving complex problems, yet their performance is critically dependent on the design of their underlying collaboration topology. As MAS become increasingly deployed in web services (e.g., search engines), designing adaptive topologies for diverse cross-domain user queries becomes essential. Current graph learning-based design methodologies often adhere to a "one-for-one" paradigm, where a specialized model is trained for each specific task domain. This approach suffers from poor generalization to unseen domains and fails to leverage shared structural knowledge across different tasks. To address this, we propose OFA-TAD, a one-for-all framework that generates adaptive collaboration graphs for any task described in natural language through a single universal model. Our approach integrates a Task-Aware Graph State Encoder (TAGSE) that filters task-relevant node information via sparse gating, and a Mixture-of-Experts (MoE) architecture that dynamically selects specialized sub-networks to drive node and edge prediction. We employ a three-stage training strategy: unconditional pre-training on canonical topologies for structural priors, large-scale conditional pre-training on LLM-generated datasets for task-topology mappings, and supervised fine-tuning on empirically validated graphs. Experiments across six diverse benchmarks show that OFA-TAD significantly outperforms specialized one-for-one models, generating highly adaptive MAS topologies. Code: https://github.com/Shiy-Li/OFA-MAS.
1 Introduction
OFA-MAS proposes a universal graph-generative framework for adaptive MAS topology design, combining task-aware encoding and expert specialization with a three-stage training curriculum and LLM-driven data synthesis. The introduction frames this approach as a response to the limits of domain-specific models and reports strong results across six benchmarks.
1 Introduction
- MAS performance depends on which agents communicate and when, while manually crafted chains, trees, and debate structures can be rigid and hard to generalize.Automatic graph-learning methods instead adapt interaction topologies to a given query.Source: §1 · p. 1
- One-for-one topology models are limited to their training domains, incur costly per-domain pipelines, and miss reusable collaboration patterns.Source: §1 · pp. 1–2
- A universal designer must encode task context and dynamically specialize expertise because fixed role pools and one-shot generation can produce suboptimal configurations.Source: §1 · p. 2
- Scarce ground-truth query–topology data and catastrophic forgetting motivate LLM-synthesized training data and an easy-to-hard curriculum.The curriculum progresses through unconditional pre-training, LLM-guided conditional training, and supervised fine-tuning.Source: §1 · p. 2
- OFA-MAS combines task-aware graph encoding, mixture-of-experts generation, three-stage training, and LLM-driven query–topology data synthesis.Its single universal designer is intended to reduce domain-specific maintenance and exploit shared cross-domain knowledge.Source: §1 · p. 2
- Experiments across six multi-domain benchmarks show OFA-MAS outperforms existing specialized methods with strong generalization.Source: §1 · p. 3
2 Problem Formulation
The section formalizes an LLM-based multi-agent system as a directed collaboration graph and contrasts domain-specific topology design with a single model trained to serve multiple domains. The one-for-all formulation targets topology generation across supported task domains without domain-specific fine-tuning or model selection.
2.1 Formulate MAS as Graphs
- Each graph vertex represents an LLM agent characterized by a functional role and a contextual state containing knowledge accumulated from prior interactions.Source: §2.1 · p. 3
- The graph’s directed acyclic edges encode information flow, and each agent’s immediate providers are the agents with edges directed into it.Source: §2.1 · p. 3
- Agents activate in topological order over communication rounds, with the system able to run for K rounds to iteratively refine messages.Source: §2.1 · p. 3
- The final system output aggregates the agents’ final-round messages using a task-appropriate combination strategy.Source: §2.1 · p. 3
2.2 From One-for-One to One-for-All Paradigm
- Traditional one-for-one design trains a domain-specific mapping from queries and a predefined role set to topologies, limiting cross-domain transferability and scalability.Source: §2.2 · p. 3
- OFA trains one designer on combined multi-domain data to generate topologies across test domains, without separate models, domain-specific fine-tuning, or model selection.Source: §2.2 · p. 3
3 Methodology
OFA-MAS generates task-adaptive multi-agent collaboration graphs with an autoregressive backbone, task-aware graph encoding, and dynamically routed experts. A three-stage curriculum moves from structural priors to task-topology alignment and empirically grounded fine-tuning.
3.1 Autoregressive Graph Generation Backbone
- OFA-MAS generates graphs stepwise by selecting an agent role and predicting its incoming edges, supporting cross-domain flexibility and variable-size topologies.The authors motivate generation from scratch as an alternative to models that fix graph scale and role allocation, which limits adaptation and generalization [43].Source: §3.1 · pp. 3–4
3.2 Task-Aware Graph State Encoder
- TAGSE conditions every graph-message-passing layer on a task vector, using sparse gating and role-aware attention to encode task-relevant information.It embeds task queries and role descriptions in a shared semantic space; the resulting node representations incorporate task context, connected-role correlations, and each node’s own state [24] [46].Source: §3.2 · pp. 4–5
3.3.1 Expert Gating Network.
- A task-vector gating network assigns weights across specialized experts, allowing different task categories to activate and combine distinct topology-design strategies.The paper illustrates sequential-chain strategies for mathematical reasoning and hierarchical review structures for code debugging.Source: §3.3 · p. 5, §3.3.1 · p. 5
3.3.2 MoE Prediction Heads.
- MoE prediction heads use expert-weighted outputs to predict the next node role and, separately, each potential incoming edge.The role head uses a pooled representation of the partial graph; the edge head combines source-node, target-node, and task representations. Both decisions use the task-specific expert ensemble.Source: §3.3.2 · p. 5
3.4.1 Stage 1: Unconditional Graph Pre-training.
- Stage 1 pre-trains the generator without task conditioning on varied classic communication graphs, establishing structural priors for valid graph construction.The dataset includes Chain, Star, and FullConnected topologies with varying sizes and role allocations; the task vector is replaced with zeros.Source: §3.4.1 · pp. 5–6
3.4.2 Stage 2: Task-Topology Alignment via LLM-driven Data Syn-
- Stage 2 uses LLM-generated task–configuration pairs for conditional pre-training, transferring broad task-topology alignment without costly MAS executions.The synthetic configurations are converted into a graph dataset and training starts from the Stage 1 model.Source: §3.4.2 · p. 6
3.4.3 Stage 3: Supervised Generative Fine-tuning.
- Stage 3 fine-tunes on task–topology pairs from real-world training queries, selecting classical configurations whose performance was validated by executing the corresponding MAS.This uses a small fraction of training data and considers structures such as sequential chains, layered topologies, and stars.Source: §3.4.3 · p. 6
3.4.4 Training Objective.
- The training objective combines teacher-forced node and edge prediction losses with expert-balancing and sparse-gating regularizers.The balance loss encourages use of all experts, while L1 regularization on TAGSE gate activations promotes selective information flow.Source: §3.4.4 · p. 6
4 Experiments
Across six benchmarks and an unseen-domain test, OFA-MAS shows strong performance, transfer, robustness, and efficiency. Ablations and analyses attribute its results to its design, pre-training, expert specialization, and adaptive role selection.
4.1 Experimental Setup
- Experiments cover MMLU, GSM8K, AQuA, MultiArith, SVAMP, and HumanEval, with GAIA used for OOD testing and success rate reported per test set.OFA-MAS is trained as one unified model across datasets, while one-for-one baselines are trained per domain; comparisons include prompting, static topologies, debate systems, and graph-learning designers.Source: §4.1 · p. 7
4.2 Main Results
- OFA-MAS achieves a state-of-the-art average score of 93.02%, outperforming all baselines by a significant margin and exceeding EIB-LEARNER by 1.64 percentage points.OFA-MAS (pre-trained) averages 92.15%, surpassing EIB-LEARNER at 91.38%.Source: §4.2 · p. 7
4.3 Out-of-Distribution Generalization
- 8.79% average accuracy: on unseen GAIA tasks, OFA-MAS outperforms all baselines and achieves 13.21% on Level-1 tasks.G-Designer averages 5.50%, below the Chain topology’s 7.38%; evaluation used GPT-4o-mini without agentic techniques.Source: §4.3 · p. 7
4.4 Ablation Study
- Ablations show that replacing tailored architectural components, removing either pre-training stage, or dropping either auxiliary loss reduces performance.The losses are balance and gate sparsity losses; the study reports qualitative drops without supplying their values in the passage.Source: §4.4 · p. 7, Table 3 · p. 7
4.5 Robustness Analysis
- OFA-MAS’s performance fell from 89.54% to 87.58% under attack, a 2.2% degradation, versus drops of 4.4% for G-Designer and 5.6% for AgentPrune.The attack replaces a key agent role with a malicious prompt intended to mislead reasoning.Source: §4.5 · pp. 7–8, Figure 3 · p. 8
4.6 Analysis of MoE Expert Specialization
- MoE gating shows consistent expert activations within domains and distinct patterns across domains, including code, general knowledge, and mathematical reasoning.HumanEval emphasizes Experts 0 and 5, MMLU activates a broader combination including Experts 6 and 7, and GSM8K and MultiArith mainly use Experts 2, 3, and 4.Source: Figure 3 · p. 8, §4.6 · p. 8
4.8 Efficiency Analysis
- On MMLU, OFA-MAS achieves 89% accuracy using 2.6e+05 tokens; on GSM8K, it achieves 95% accuracy at the lowest token cost, 6.3e+06.The passage describes its accuracy as superior with competitive efficiency on MMLU and GSM8K.Source: Figure 4 · p. 8, §4.8 · p. 8
4.7 Case Study
- For HumanEval, OFA-MAS forms a sequential refinement chain with code-generation roles; for GSM8K, it combines math-relevant roles with a Critic role associated with MMLU.The examples illustrate role selection from a universal pool across task types.Source: §4.7 · p. 8
5 Conclusion … A.2 Autoregressive Graph Generation
OFA-MAS combines task-aware graph encoding, MoE specialization, and three-stage training to generate adaptive MAS collaboration topologies. The paper reports that this universal model generalizes across diverse domains and achieves state-of-the-art performance.
5 Conclusion
- Extensive experiments show that OFA-MAS generalizes across diverse domains and achieves state-of-the-art performance by leveraging shared knowledge.The conclusion presents this as evidence that a universal model can overcome limitations of specialized one-for-one designs.Source: §5 · p. 8
- OFA-MAS combines a task-aware graph state encoder (TAGSE), a mixture-of-experts (MoE) architecture, and three-stage training to generate task-specific topologies from general collaborative-graph knowledge.The conclusion frames this progression as a path from specialized solutions toward a universal graph foundation model for MAS design.Source: §5 · p. 8
A.1 Multi-Agent Systems Topology Design
- Early topology designs used fixed, domain-specific templates—sequential chains [35] [44], tree-of-thought hierarchies [1] [37], and debate structures [7] [11] [19]—that lack adaptability across task requirements.These approaches followed earlier reinforcement learning-based coordination [32] and later LLM-based collaboration frameworks.Source: §A.1 · p. 10
- Learnable approaches include pruning dense graphs with AgentPrune [42] and Agent-Drop [34], direct graph generation with GPTSwarm and G-Designer [43], and evolutionary methods [10] [40] and neural architecture search [41].The paper says these methods generally train specialized models for individual domains, limiting scalability and cross-domain knowledge transfer; OFA-MAS instead uses one-for-all MoE specialization.Source: §A.1 · p. 10
A.2 Autoregressive Graph Generation
- Autoregressive graph generation became a dominant paradigm for its compatibility with sequential modeling: GraphRNN [39] generates nodes and edges sequentially, while later work improved graph representations and neural architectures [15].Source: §A.2 · p. 10
- PARD [45] uses permutation-invariant autoregressive diffusion to address ordering sensitivity, G2PT [4] applies transformers with strong generalization, and Orderless Regularization [6] encourages invariance across valid graph orderings.Source: §A.2 · p. 10
B Dataset Details · B.1 In-Distribution Benchmarks
The in-distribution evaluation spans six benchmarks covering broad knowledge, mathematical reasoning, and code generation, with benchmark-specific test subsets and task formats.
B.1 In-Distribution Benchmarks
- MMLU covers 57 subjects; OFA-MAS uses 153 test questions, each with four answer options.The subjects span STEM, humanities, social sciences, and more.Source: §B.1 · p. 10
- GSM8K contributes 1,000 test-split problems from a dataset of 8.5K grade-school math word problems.Each problem requires multi-step arithmetic reasoning.Source: §B.1 · p. 10
- AQuA includes 254 algebraic word problems with natural-language rationales and five answer options per question.The problems require multi-step algebraic reasoning.Source: §B.1 · p. 10
- MultiArith contains 600 elementary-level arithmetic word problems requiring multi-step reasoning and numerical answers.Source: §B.1 · p. 10
- SVAMP uses 1,000 problems with subtle linguistic variations to test robustness and model generalization.Source: §B.1 · p. 10
- HumanEval contains 164 code-generation problems and is evaluated using Pass@1.Each problem includes a function signature, docstring, body, and unit tests.Source: §B.1 · p. 10
B.2 Out-of-Distribution Benchmark
GAIA tests OFA-MAS on an unseen, knowledge-intensive benchmark requiring multimodal reasoning, tool use, and web search. OFA-MAS is evaluated zero-shot, while one-for-one baselines use models trained on MMLU.
B.2 Out-of-Distribution Benchmark
- GAIA was never seen during training and tests real-world questions requiring multimodal reasoning, tool use, and web search.It includes 53 Level-1, 86 Level-2, and 26 Level-3 questions, with levels differing in complexity and reasoning requirements.Source: §B.2 · p. 10
- OFA-MAS is applied to GAIA without additional training, testing zero-shot OOD generalization from its six in-distribution benchmarks.The unified model is evaluated directly on GAIA.Source: §B.2 · p. 10
- One-for-one baselines use models trained on MMLU, selected for GAIA because of its knowledge-intensive, multimodal reasoning demands.The protocol provides baselines with the most relevant role configurations and domain knowledge for the OOD task.Source: §B.2 · p. 10
B.3 Data Statistics · C Baseline Details
This block describes the scope of the dataset statistics and the single-agent, fixed-topology, and learning-based baselines used in the experiments.
B.3 Data Statistics
- Table 4 summarizes dataset answer types, evaluation metrics, and licensing information.The table is identified as providing comprehensive statistics for all experimental datasets.Source: §B.3 · p. 10, Table 4 · p. 11
C Baseline Details
- Single-agent baselines are Vanilla direct prompting, CoT [35] with step-by-step intermediate reasoning, and SC [33], which majority-votes across five CoT paths.Source: C Baseline Details · p. 11
- Fixed topologies include Chain’s sequential links, Tree’s hierarchy, and Complete’s fully connected graph for maximum information sharing.Source: C Baseline Details · p. 11
- The Random baseline uses randomly generated topologies and role assignments, while LLM-Debate has agents critique one another over multiple rounds.Source: C Baseline Details · p. 11
- AgentPrune [42] performs one-shot pruning on spatial-temporal message-passing graphs to reduce communication redundancy and token costs while defending against adversarial attacks.Source: C Baseline Details · p. 11
- AgentDropout [34] optimizes adjacency matrices to remove redundant agents and communications across rounds, aiming to improve token efficiency and task performance.Source: C Baseline Details · p. 11
- G-Designer [43] uses a variational graph auto-encoder with agent and task-specific virtual nodes to design task-adaptive communication topologies.Source: C Baseline Details · p. 11
- EIB-LEARNER [30] uses causal error-propagation analysis and combines dense and sparse connectivity patterns to balance error suppression with information diffusion.Source: C Baseline Details · p. 11
D Implementation Details · D.1 Model Architecture
OFA-TAD combines a task encoder, a task-aware graph state encoder, and a mixture-of-experts module with a universal role pool for autoregressive generation.
D.1 Model Architecture
- The task encoder uses all-MiniLM-L6-v2 to produce 384-dimensional sentence embeddings, then projects them to 128 dimensions through a two-layer MLP with hidden dimension 256.Source: §D.1 · p. 11
- TAGSE uses four layers of context-gated message passing with role-aware attention aggregation and hidden dimension 256.Source: §D.1 · p. 11
- TAGSE initially projects node features from 384 dimensions to 256 dimensions.Source: §D.1 · p. 11
- The MoE module has eight expert networks, each a two-layer MLP with hidden dimension 256 and ReLU activation.Source: §D.1 · p. 11
- A two-layer gating MLP maps 128 dimensions through 256 to eight outputs, using ReLU and Softmax.Source: §D.1 · p. 11
- ROFA contains 19 diverse roles spanning multiple domains, alongside START and END tokens for autoregressive generation.Source: §D.1 · p. 11
- ROFA’s 19-role pool is larger than the 4–6 roles per domain typically used by one-for-one methods.Source: §D.1 · p. 11
D.2 Training Configuration … E LLM Prompt Templates
The appendix specifies OFA-TAD’s staged training, loss and inference settings, computational environment, and LLM prompts for synthesizing task-topology pairs and defining agent roles.
D.2 Training Configuration
- Training uses 800 canonical graphs for 50-epoch unconditional pre-training, then 1,000 GPT-4o-mini-generated pairs for conditional training and 1,000 empirically validated pairs for fine-tuning.The first two stages use batch size 32, learning rate 2 × 10−3, Adam, and ReduceLROnPlateau (factor 0.5, patience 10); fine-tuning initializes from Stage 2 and covers 40 queries per benchmark across six benchmarks.Source: §D.2 · p. 11
- Classical topology candidates are Chain, Layered, Star, Mesh, and FullConnected, with the top two per query selected by validation accuracy.This training configuration uses 50 epochs, batch size 32, learning rate 5 × 10−4, Adam, and ReduceLROnPlateau (factor 0.5, patience 10).Source: §D.2 · p. 12
E.1 Prompt for LLM-Driven Data Synthesis (Stage 2)
- The synthesis instructions balance agent counts from 2–6, use all five topologies and 19 roles, require directed acyclic graphs, and avoid redundant topology choices for smaller teams.For two agents, only Chain is used; for three, FullConnected is excluded as identical to Mesh; for four or more, all five topologies are distinct.Source: §D.2 · p. 12
D.3 Loss Function Hyperparameters
- The loss combines graph loss with MoE balance loss weighted 0.2 and gate sparsity L1 loss weighted 0.1.L = Lgraph + 𝜆balanceLbalance + 𝜆gateLgate.Source: §D.3 · p. 12
D.4 Inference Configuration
- Inference limits generated graphs to 2–6 nodes and uses greedy decoding or multinomial sampling, stopping at END and sampling edges probabilistically.If no edges are sampled, the system forces a connection to the highest-probability node; MAS execution uses K = 1 communication round with GPT-4o agents at temperature 0.0.Source: §D.4 · p. 12
D.5 Computational Resources
- Experiments ran on a machine with a 12-vCPU Intel Xeon Platinum 8255C, one RTX 2080 Ti with 11 GB GPU memory, Ubuntu 20.04, Python 3.8, PyTorch 2.1.2, and CUDA 12.1.Source: §D.5 · p. 12
E.1 Prompt for LLM-Driven Data Synthesis (Stage 2)
- Stage 2 synthesis prompts use GPT-4o at temperature 0.7 to generate diverse queries and suitable MAS configurations across domains, roles, and five topology types.The illustrated JSONL-style example includes a query, agent count, roles, topology, and edges.Source: §E.1 · p. 12
E.2 Agent Role Prompts
- The universal role pool contains 19 specialized roles, with prompts covering knowledge-entity search, mathematics, and concise algorithm design.The complete-pool list includes roles such as Critic, Doctor, Project Manager, Programming Expert, and Bug Fixer, plus domain-specific variants for choice questions.Source: §E.2 · p. 12