MoRAgent: Parameter Efficient Agent Tuning with Mixture-of-Roles

AI summary

Overview

• Research area: Natural Language Processing, specifically parameter-efficient fine-tuning (PEFT) of large language models for agent tasks such as function calling and tool use.
• Technical level: Advanced. The paper assumes familiarity with LoRA, mixture-of-experts routing, low-rank decomposition, and agent benchmarks.
• Scope in one sentence: The paper introduces MoRAgent, a method that splits agent capabilities into three roles (reasoner, executor, summarizer) and trains each with its own group of LoRA adapters, achieving competitive agent performance while tuning only a small fraction of parameters.

What This Paper Is About

Fine-tuning open-source LLMs for agent tasks usually means full-parameter fine-tuning, which is expensive and degrades the base model's general abilities. Parameter-efficient methods like LoRA are cheap but underperform on agents, because a low-rank adapter must simultaneously learn reasoning, function execution, and summarization. The paper's goal is to close that gap by decomposing agent capability into distinct roles and giving each role its own dedicated set of LoRA modules under a Mixture-of-Roles framework.

Key Contributions

1. Capability decomposition into three roles. Inspired by the dominant Reason+Action paradigm, the authors split agent capability into a reasoner (understands the query and decides the next role), an executor (selects functions and parameters), and a summarizer (distills the conversation back to the user). The summarizer is only invoked when the reasoner judges the query solved, or when the query cannot be solved after enough executor attempts.

2. Mixture-of-Roles (MoR) framework. Three specialized LoRA groups are attached to the frozen backbone, one per role, with each group containing a shared LoRA plus a number of routed LoRAs. A rule-based role-aware gate ensures only one role is active per token, while learnable Top-K token-aware routers choose which LoRAs within a role handle a given input. The framework can be placed on the linear layers of attention or the feed-forward network in each transformer block.

3. New training objectives. Alongside cross-entropy loss, the paper adds an auxiliary balance loss (inherited from Switch Transformers) to avoid load imbalance across LoRAs, and an orthogonal loss between LoRA A and B matrices to push different LoRAs toward different feature directions.

4. Multi-role data generation pipeline. Built on public datasets (ToolBench, APIGen, ToolACE, glaive-function-calling-v2, MathGenie), the pipeline uses GPT-4o to complete missing role content such as thoughts and summaries, DeepSeek-V3 to score trajectory and step quality, and a hybrid of rule-based filtering, LLM prompting, and manual correction to fix executor errors.

Main Findings

• StableToolBench gains: On Llama3.2-1B-Instruct (1.24B), the paper reports a DFS pass-rate improvement of 44.5% and a win-rate improvement of 25.2% over the base model, using 0.16B additional trainable parameters. On Phi3.5-mini-Instruct (3.82B), the reported CoT improvements are 26.3% pass rate and 17.3% win rate with 0.36B trainable parameters. Note that the abstract states DFS pass-rate improvements of 40.6% and 14.2% for these two models, which does not match the 44.5% shown for Llama in Table 1 and Section 4.3; the 14.2% figure matches the Phi DFS result in Table 1.

• ToolLLaMA-v2 comparison: MoRAgent-Phi, with fewer parameters, surpasses the agent-specific ToolLLaMA-v2-7B on StableToolBench.

• BFCL leaderboard: Average accuracy improves by 50.1% for Llama3.2-1B-Instruct (27.5 to 77.6, 0.16B trainable), 17.9% for Phi3.5-mini-Instruct (64.2 to 82.1, 0.36B), and 53.4% for openPangu-Embedded-1B (23.8 to 77.2, 0.20B). For reference, the paper lists Qwen2.5-72B-it at 89.7, GPT4 at 85.2, ToolACE-8B at 84.6, MiniCPM3-4B at 77.7, and Llama3.2-3B-it at 77.3.

• Math tasks: Using Qwen2.5-1.5B-Coder (1.54B) and solving problems by importing and executing Python packages, MoRAgent-Qwen reaches 68.5 on GSM8K (+13.8) and 45.5 on MATH (+12.0) with 0.27B trainable parameters, versus the best baseline of 54.7 (GSM8K, with packages) and 33.5 (MATH, without packages).

• Loss ablation: On BFCL with Llama3.2-1B-Instruct, cross-entropy alone gives 72.3 average accuracy, adding balance loss gives 75.1, adding orthogonal loss gives 74.4, and combining all three gives 77.6, a 5.3% improvement. Similarity visualizations of the query module at layer 7 show routed LoRAs become less similar with the orthogonal loss.

• LoRA count ablation: Increasing LoRAs raises accuracy but also parameter count: 4/4/3 (reasoner/executor/summarizer) gives 0.13B and 73.2; 5/5/4 gives 0.16B and 77.6; 6/6/5 gives 0.19B and 79.1; 7/7/6 gives 0.23B and 80.0.

• Method comparison: With the same multi-role dataset, LoRA scores 65.5, DoRA 66.8, full-parameter SFT 76.0 (at 1.24B trainable), and MoRAgent 77.6 at 0.16B trainable parameters.

• Data efficiency: Even with only 1K training samples, average BFCL accuracy improves by 27.4% (27.5 to 54.9). Accuracy rises to 60.0 at 5K, 63.7 at 10K, 73.4 at 50K, and 77.6 at 90K.

• Role difficulty differs: Toy experiments (80K training and 5K validation samples per role, rank 16, 2 epochs on Llama3.2-1B-Instruct) show loss and Levenshtein accuracy stop improving sharply after 4 to 5 LoRAs for the reasoner and executor, but after 3 to 4 for the summarizer, which the authors interpret as the summarizer being an easier role to learn.

Methodology in Plain English

The authors start from a simple observation: a single low-rank adapter has to learn too many things at once. So they first define what an agent model actually does, in three separate jobs. The reasoner reads the user query and the history, writes an analysis, and announces which role should act next. The executor reads that analysis and picks the function name and parameters to call. The summarizer takes the whole conversation and writes the final answer for the user.

Then they build a training architecture to match. The base model's weights stay frozen. On top of certain linear layers, they attach three separate pools of LoRA modules, one pool per role. A rule-based gate makes sure that for any given token, only one role's LoRA pool is active, so roles do not interfere with each other. Inside each pool, a small learnable router picks the top few LoRAs to apply, plus one shared LoRA that always applies. The number of LoRAs differs by role: five total (four active) for the reasoner and executor, four total (three active) for the summarizer, based on the toy experiments showing the summarizer converges faster.

Training uses three losses at once: standard cross-entropy for next-token prediction, a balance loss that encourages tokens to spread across the LoRAs in a pool rather than collapsing onto one, and an orthogonal loss that makes different LoRAs' matrices point in different directions so they capture complementary features. Hyperparameters are a rank of 16, learning rate 5e-5, 4 epochs, and loss weights of 1e-3 and 1e-4. The seven target modules tuned are query, key, value, out, gate, up, and down.

The data side matters as much as the architecture. Public agent datasets often give only the action, not the reasoning or summary. The authors use GPT-4o to fill in those missing role contents, DeepSeek-V3 to grade the completed trajectories, and then manually or automatically repair executor mistakes such as calling functions not in the candidate list, using wrong parameter counts or types, or choosing the wrong function in a way that does not crash. Everything is converted to a unified JSON format for fine-tuning.

Why This Matters

• Research impact: The paper shows that how you allocate a fixed parameter budget across an agent's sub-capabilities may matter more than the budget itself. Its method beats full-parameter SFT on BFCL average accuracy (77.6 versus 76.0) while training roughly an eighth as many parameters, which challenges the assumption that full fine-tuning is the ceiling for agents.

• Real-world applications:

  • On-device or edge assistants where only small models like Llama3.2-1B-Instruct or openPangu-Embedded-1B can run, and where the reported 0.16B to 0.20B trainable parameter overhead is affordable.
  • Tool and API assistants that must pick correct functions and arguments, evaluated here on StableToolBench's 765 questions across 6 subtasks and BFCL's 2797 questions.
  • Math and quantitative reasoning assistants that solve problems by generating and executing Python, as done with MathGenie data on GSM8K and MATH.
  • Multi-turn customer-facing agents that need a distinct final summarization step to convey results back to users in natural language.

• Industry relevance: The two stated motivations for avoiding full fine-tuning are the compute cost of billions of parameters and the loss of a base model's general capabilities. MoRAgent's frozen backbone and small adapter footprint directly address both, which matters for organizations that want to switch the same base model between general chat and agent workloads.

Future Directions

• Resolve the reporting discrepancy: The abstract's 40.6% DFS pass-rate improvement for Llama3.2-1B-Instruct does not correspond to the 44.5% in Table 1 and Section 4.3, so the source of that number is not explained in the paper.

• Scale beyond small models: All reported base models are 3.82B parameters or smaller (Llama3.2-1B-Instruct, Phi3.5-mini-Instruct, Qwen2.5-1.5B-Coder, openPangu-Embedded-1B). The paper does not report whether the role decomposition and per-role LoRA allocation logic holds at larger scales.

• Generalize the role decomposition: The three-role split is fixed by hand. Whether other agent settings need more roles, fewer roles, or learned roles is not explored in the paper.

• Combine with other PEFT advances: The paper compares against LoRA and DoRA but does not report combinations with quantized or other adapter variants, nor does it report whether the routers transfer across tasks without retraining.

Target Audience

This paper is most useful to researchers and engineers working on parameter-efficient fine-tuning, LLM agents, and function-calling systems. Practitioners who deploy small open-source models as tool-using agents under tight compute budgets will find the concrete numbers and the reasoner/executor/summarizer design directly actionable. Readers without a background in LoRA, MoE routing, or agent benchmarks will need to consult the cited PEFT and agent-tuning literature first, since the method section assumes that context.

Authors’ abstract

Despite recent advancements of fine-tuning large language models (LLMs) to facilitate agent tasks, parameter-efficient fine-tuning (PEFT) methodologies for agent remain largely unexplored. In this paper, we introduce three key strategies for PEFT in agent tasks: 1) Inspired by the increasingly dominant Reason+Action paradigm, we first decompose the capabilities necessary for the agent tasks into three distinct roles: reasoner, executor, and summarizer. The reasoner is responsible for comprehending the user's query and determining the next role based on the execution trajectory. The executor is tasked with identifying the appropriate functions and parameters to invoke. The summarizer conveys the distilled information from conversations back to the user. 2) We then propose the Mixture-of-Roles (MoR) framework, which comprises three specialized Low-Rank Adaptation (LoRA) groups, each designated to fulfill a distinct role. By focusing on their respective specialized capabilities and engaging in collaborative interactions, these LoRAs collectively accomplish the agent task. 3) To effectively fine-tune the framework, we develop a multi-role data generation pipeline based on publicly available datasets, incorporating role-specific content completion and reliability verification. We conduct extensive experiments and thorough ablation studies on various LLMs and agent benchmarks, demonstrating the effectiveness of the proposed method. This project is publicly available at https://mor-agent.github.io.