💻 computer science

HieraMAS: Optimizing Intra-Node LLM Mixtures and Inter-Node Topology for Multi-Agent Systems

HieraMAS is a hierarchical multi-agent framework that optimizes both intra-node LLM mixtures and inter-node communication topologies through a two-stage algorithm addressing credit-assignment challenges, achieving superior performance and cost-efficiency on reasoning and coding benchmarks.

Original authors: Tianjun Yao, Zhaoyi Li, Zhiqiang Shen

Published 2026-02-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tianjun Yao, Zhaoyi Li, Zhiqiang Shen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a very difficult puzzle, like a complex math problem or writing a piece of software. You could hire one genius to do it all, but that's expensive and risky if they get tired or stuck. Alternatively, you could hire a team of people.

Most current "AI teams" work like a rigid corporate structure: you assign a specific job to a specific person (e.g., "You are the Coder," "You are the Editor"), and they talk to each other in a fixed pattern (like a chain or a circle).

HieraMAS is a new, smarter way to organize these AI teams. Think of it as a dynamic, self-optimizing startup rather than a rigid corporation. Here is how it works, broken down into simple concepts:

1. The "Super-Node" (The Internal Brainstorming Session)

In traditional AI teams, one "agent" is just one AI model. In HieraMAS, every role is actually a Super-Node.

The Analogy: Imagine you need a "Math Expert." Instead of hiring just one mathematician, you hire a small committee inside that role.
- The Proposers: You have three different mathematicians (maybe one is fast but sloppy, one is slow but precise, one is great at geometry). They all write down their own solutions independently.
- The Synthesizer: A fourth person (the editor) reads all three proposals, picks the best parts, and writes the final answer.
The Magic: HieraMAS learns which combination of mathematicians works best for which specific problem. Sometimes it needs the fast one; sometimes it needs the precise one. It can even decide to fire (skip) a mathematician if they aren't needed, saving money.

2. The "Dynamic Org Chart" (The Communication Topology)

Once you have your Super-Nodes, they need to talk to each other. Traditional systems use a fixed map: "The Coder talks to the Editor, who talks to the Manager."

The Analogy: HieraMAS doesn't use a fixed org chart. It uses a dynamic network.
- For a coding task, the "Coder" might talk directly to the "Bug Fixer."
- For a history task, the "Historian" might talk to the "Lawyer" and the "Economist," skipping the Coder entirely.
The Magic: The system learns the best "conversation path" for every single question. It builds a custom map for every task, ensuring information flows exactly where it's needed, no more, no less.

3. The Two Big Problems (And How They Solved Them)

The authors realized that building this flexible system was incredibly hard because of two "blame games":

Problem A: The "Good Teammate" Effect (Node Credit Assignment)

The Scenario: Imagine a team where one person makes a huge mistake, but another person fixes it so well that the final result is perfect.
The Mistake: If you only look at the final grade (A+), you might think the person who made the mistake did a great job! You accidentally reward bad behavior.
The HieraMAS Fix: They use Multi-Level Rewards. It's like a teacher who gives a grade for the final essay and a separate grade for each student's draft. Even if the final essay is perfect, the student who wrote the bad draft gets a "bad grade" for their specific part. This ensures everyone is held accountable for their own work.

Problem B: The "Who Said What?" Effect (Edge Credit Assignment)

The Scenario: In a complex network, it's impossible to tell if a specific conversation (an "edge" between two people) helped or hurt the result. Was it the link between the Coder and the Editor that saved the day, or the link between the Editor and the Manager?
The Mistake: Trying to tweak every single connection one by one is like trying to tune a radio by turning every single screw on the circuit board. It's too messy and confusing.
The HieraMAS Fix: They stopped trying to tune individual wires. Instead, they treated the whole network as a single object.
- The Analogy: Imagine you have 100 different organizational charts (some are circles, some are trees, some are webs). Instead of trying to fix one line on one chart, you run the task 100 times with 100 different charts, see which chart got the best result, and say, "Okay, this whole chart is the winner." It's much easier to pick the best whole picture than to fix the pixels one by one.

The Result

By combining internal mix-and-match teams (Super-Nodes) with dynamic communication maps, HieraMAS achieves two things:

Smarter Results: It solves hard problems better than existing methods because it uses the right tools for the right job.
Cheaper Costs: It doesn't waste money using expensive, super-smart AI models for simple tasks, and it doesn't waste time having agents talk to each other when they don't need to.

In short: HieraMAS is like a master chef who doesn't just hire one sous-chef, but builds a custom kitchen crew and a custom workflow for every single dish, ensuring the meal is delicious without burning the budget.

1. Problem Statement

Existing Multi-Agent Systems (MAS) built on Large Language Models (LLMs) typically treat each agent as a monolithic unit. Current optimization efforts focus on isolated dimensions:

Communication Topology: Optimizing how agents connect (e.g., chain, star, fully connected).
Role Assignment: Determining which agent performs which function.
LLM Routing: Selecting a single backbone model for a specific role.

The Gap: These approaches fail to exploit internal LLM mixtures. Recent research suggests that LLMs perform better when provided with outputs from other models as auxiliary input (the "collaborativeness" phenomenon). However, current MAS frameworks do not integrate this intra-agent collaboration with inter-agent topology optimization.

The Core Challenge: Jointly optimizing intra-node configurations (LLM mixtures) and inter-node structures (topology) creates severe credit assignment problems:

Node-Level Masking: A final task reward might be high because other agents compensated for a specific node's error, leading to incorrect reinforcement of a suboptimal node configuration.
Edge-Level Entanglement: Optimizing individual communication edges is difficult because the contribution of a single edge is entangled with the behaviors of both sending and receiving nodes, making it hard to isolate the edge's specific impact.

2. Methodology: HieraMAS

HieraMAS introduces a hierarchical framework that unifies intra-node and inter-node collaboration.

A. Core Architecture: Supernodes

Instead of single agents, HieraMAS uses Supernodes.

Structure: Each functional role (e.g., "Math Solver") is a supernode containing a Propose-Synthesis structure.
- Proposers ( $W$ ): Multiple heterogeneous LLMs generate diverse proposals.
- Synthesizer: A single LLM aggregates these proposals into a unified response.
Flexibility: The system includes a special "skip" token.
- If selected for a proposer, that proposer is omitted (cost reduction).
- If selected for a synthesizer, the entire supernode (role) is deactivated (role pruning).

B. Optimization Dimensions

HieraMAS jointly optimizes three interconnected dimensions:

Graph Topology: The communication pattern between supernodes.
Role Pruning: Deciding which functional roles are necessary for a specific query.
LLM Selection: Choosing the optimal mix of LLMs (from a heterogeneous pool) for the proposer and synthesizer positions within each supernode.

C. Two-Stage Training Algorithm

To solve the credit assignment challenges, HieraMAS employs a decoupled, two-stage training process:

Stage 1: Supernode Optimization (with Random Graphs)

Goal: Learn optimal LLM selection and role pruning.
Mechanism: The system executes the MAS using randomly sampled communication topologies from a pre-generated pool. This prevents overfitting to a specific graph structure.
Solution to Challenge 1 (Node Masking): Uses Multi-Level Rewards.
- Instead of relying solely on the final task reward ( $R_{final}$ ), the system computes a node-level reward ( $R_{node}$ ) based on the synthesizer's output quality.
- Effective Reward: $R_{eff} = \alpha \cdot R_{node} + (1-\alpha) \cdot R_{final}$ .
- This ensures that even if the final task succeeds due to other agents, a failing node receives a negative signal proportional to its individual error, preventing erroneous reinforcement.

Stage 2: Graph Classifier Training

Goal: Learn the optimal communication topology.
Mechanism: The LLM selector is frozen. The system treats topology selection as a holistic graph classification problem rather than per-edge optimization.
Process:
1. A pool of $K$ diverse Directed Acyclic Graphs (DAGs) is pre-generated.
2. The system executes the MAS with different graphs from the pool using the frozen LLM selector.
3. A Graph Classifier (based on a Graph Convolutional Network) is trained to score these candidate graphs and select the best one for a given query.
Solution to Challenge 2 (Edge Entanglement): By evaluating the entire topology as an indivisible unit, the method bypasses the intractable problem of assigning credit to specific edges.

3. Key Contributions

Novel Paradigm: Introduced a two-level collaboration framework: Intra-node (heterogeneous LLM mixtures within supernodes) and Inter-node (communication across functional roles).
Unified Framework: HieraMAS is the first to jointly optimize topology, role retention, and LLM selection within a single coherent system.
Credit Assignment Solutions:
- Proposed Multi-Level Rewards to resolve per-node attribution errors.
- Reformulated topology optimization as a Graph Classification task to circumvent per-edge credit assignment ambiguity.
Theoretical Justification: Provided proofs showing that multi-level rewards correct gradient directions and that holistic graph selection reduces estimation error to a vanishing rate, unlike per-edge optimization which suffers from irreducible error.

4. Experimental Results

The authors evaluated HieraMAS on three benchmarks: HumanEval++ (coding), MATH (reasoning), and MMLU-Redux (general knowledge).

Performance: HieraMAS achieved State-of-the-Art (SOTA) performance with an average accuracy of 94.61%, outperforming strong baselines like AFlow, GDesigner, and MASRouter.
- It achieved the best results on HumanEval++ and MATH.
Cost-Efficiency:
- HieraMAS significantly reduced inference costs compared to "Full-Graph" baselines (e.g., 3.27x cheaper on MMLU-Redux).
- It achieved better cost-performance trade-offs than using the strongest LLM for all nodes by adaptively mixing weaker (cheaper) and stronger models based on task requirements.
Ablation Studies:
- Removing the graph scoring mechanism caused a ~3.4% performance drop, highlighting the importance of learned topology.
- Removing LLM selection (using only the strongest model) increased costs by >200% while often lowering accuracy.
Generalization: HieraMAS demonstrated strong out-of-domain generalization on unseen MMLU subjects, outperforming other learning-based methods like MASRouter.

5. Significance

HieraMAS represents a significant shift in Multi-Agent System design. By moving away from monolithic agents and fixed topologies, it leverages the collaborative potential of heterogeneous LLMs at both the micro (intra-node) and macro (inter-node) levels.

The framework's ability to dynamically prune unnecessary roles and select cost-effective LLMs while maintaining high performance offers a practical solution for deploying scalable, efficient, and robust AI agents in real-world applications. The proposed two-stage training algorithm provides a robust blueprint for solving complex credit assignment problems in joint optimization settings.