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

Dear Reviewer Mktg:

Grateful for your support and helpful review. All concerns and questions are meticulously responded.

1. Experiments with a single role is supposed to be conducted.

Thanks for the helpful suggestion. We supplement the experiment of a single lora without roles on BFCL leaderboard with Llama3.2-1B-Instruct.

With the same number of trainable parameters, due to the lack of decomposition capability, the accuracy improvement achieved by LoRA is much lower than that of MoR.

2. The work of data preparation and unification is essential, but seems time-consuming and labor-intensive.

Thanks for the kind concern. Actually, a large proportion of the raw training data lacks multi-roles content, such as the thought of reasoner. Therefore, it cannot be directly used for MoR training. The data processing pipeline proposed in the paper addresses this by generating high-quality multi-role data based on the raw dataset. However, as the reasoner paradigm like DeepSeek-R1 and OpenAI o3, etc., has become mainstream recently, there will be more agent related data with thought.

3. Additional experiments on larger models.

Thanks for the helpful suggestion. Due to the limited time of rebuttal, we only supplement the experiment of Llama3.1-8B-Instruct on BFCL.

The experimental settings are consistent with section4.2 in our paper. From the results, with the introduction of 0.59B trainable parameters, we improve the average accuracy by 11.1%.

4. In table 3, why qwen performs better after training when the base results are worse than llama on MATH benchmarks.

The possible reason is that llama is a general model and qwen-coder is a code-specific model. Since we solve math problems by importing packages here, it is more in line with the style of qwen-coder, so the accuracy of qwen-coder can be improved higher in the post-training stage.

5. Using more or fewer roles have a greater impact on the results?

Thanks for the constructive suggestion, we supplement the experiments of more or fewer roles. The results are as follows.

Our method integrates three core roles: Reasoner, Executor, and Summarizer. Through removing the Summarizer and introducing a Planner, we extend the roles to 2-roles and 4-roles. Experimental results demonstrate accuracy improvements across all variants compared to the baseline model. Notably, the 3-roles architecture achieves peak performance (accuracy improvement of 50.1% over baseline), whereas the 4-roles, despite containing more trainable parameters, shows diminishing returns (45.7% accuracy gain). This empirical evidence highlights that ​rational role definition outweigh mere parameter quantity expansion in achieving optimal model performance.

6. Typos.

Sincerely thanks for your detailed reviews. We will go through the entire paper again and fix potential typos in the revised version.

7. The code of the architecture and training process are suggested to be released.

Thanks, all our data and code will be open-sourced to contribute to the community.

Dear Reviewer MDgP:

We deeply appreciate your support and insightful feedback, detailed rebuttals to all queries are provided.

1. The difference between ours and α-UMi.

α-UMi decomposes the agent ability into planner,executor and summarizer. However, each role is implemented by a separate LLM, which results in a significant increase in computing resources. In contrast, we integrate the multiple agent capabilities into a novel parameter-efficient Mixture-of-Roles framework.

2. The auxiliary balance loss is not novel.

Thanks for the kind concern. Auxiliary balance loss is widely adopted in training MoE architectures. We borrow the idea of ​​balance loss and are the first to apply it to our agent work. The experimental results on various benchmark verify the effectiveness of the loss.

3. Ablation on the auxiliary balance loss and orthogonal loss weight coefficients.

Thanks for the helpful advice. we supplement the ablation studies on BFCL leaderboard with Llama3.2-1B-Instruct.

With appropiate coefficients of auxiliary balance loss and orthogonal loss, the accuracy may be further improved.

4. The preparation of experimental data seems to be a lot of work.

Actually, a large proportion of the raw training data lacks multi-role content, such as the thought of reasoner. Therefore, it cannot be directly used for MoR training. The data processing pipeline proposed in the paper addresses this by generating high-quality multi-role data based on the raw dataset. However, as the reasoner paradigm has become mainstream recently, there will be more agent related data with thought. All data and code will be open-sourced to contribute to the community.

5. Will the two gates lead to an increase in FLOPs or latency?

Thanks. Our gates include the rule-based role-aware gate and learnable token-aware gate. For the rule-based gate, a "role" tag is designed to indicate which role activates the subsequent tokens. Therefore, no FLOPs or latency is introduced. For each learnable token-aware gate, we implement it with only a single Linear layer, resulting in a negligible FLOPs and latency.

6. The result of having three separate LLMs to represent different roles.

This is an interesting question, and actually, that is precisely how α-UMi operates. We supplement this experiment on BFCL leaderboard with Llama3.2-1B-Instruct.

With more trainable parameters, three separate LLMs achieve higher accuracy, but the computational resources required for training and inference are much higher than ours.

7. The comparison between a large LoRA and multiple small LoRAs?

Sorry for not fully understanding what multiple small LoRAs meant, we guess it should be our proposed MoR. We supplement the experiment of one large LoRA with Llama3.2-1B-Instruct on BFCL.

The LoRA is trained on the same 90k multi-roles data. From the results, the proposed MoR can fully utilize the limited training parameters and achieve significantly better accuracy.

8. Typos.

Sincerely thanks for your detailed reviews. We will go through the entire paper again and fix potential typos in the revised version.

Dear Reviewer rpMC:

We sincerely thanks for your support and meticulous review. The concerns and questions are answered as follows.

1. How is the rule-aware gate implemented during training? Are there labels in the training data indicating which role should be active for each token?

Yes! Your deduction is correct. Our training data includes a "role" tag to indicate which role activates the subsequent tokens. Take the Math data in Figure 3 in our paper as an example. The tokens marked as "content" are activated by the role marked as "role". It should be noted that when the role is user or observation, it is also activated by reasoner. Therefore, "system_prompt", "user_query", "thought and next role" and "code execution result" are activated by reasoner. "Python code" is activated by executor. "summarize and feed back to user" is activated by summarizer. We will further clarify this in the revised version.

2. The effect of multi-role data generation pipeline is unclear. It is unclear whether the performance is coming from MoR or the data pipeline. How's the performance for standard LoRA?

Thanks for the helpful suggestion. A large proportion of the raw training data lacks multi-role content, such as the thought of reasoner. Therefore, it cannot be directly used for MoR training. The data processing pipeline proposed in the paper addresses this by generating high-quality multi-role data based on the raw dataset. A comparative example is illustrated below.

Raw data:

{
    "role": "user"
    "content": "I want to watch a movie tonight. Can you suggest me some movies starring Leonardo DiCaprio?"
},
{
    "role": "assistant"
    "content": "[unused11]Action: search_movies\n Arguments: {"actor": "Leonardo DiCaprio"}[unused12]"
},
{
    "role": "observation"
    "content": "{"movies": ["Inception", "The Wolf of Wall Street", "Titanic", "The Revenant", "Shutter Island"]}"
},
{
    "role": "assistant",
    "content": " "Inception", "The Wolf of Wall Street", "Titanic", "The Revenant", "Shutter Island". "
}

Multi-role data after generation pipeline:

{
    "role": "user",
    "content": "I want to watch a movie tonight. Can you suggest me some movies starring Leonardo DiCaprio?"
},
{
    "role": "reasoner",
    "content": "To suggest movies starring Leonardo DiCaprio, we can use the "search_movies" tool with the input specifying the actor's name.\n Next: executor."
},
{
    "role": "executor",
    "content": "[unused11]Action: search_movies\n Arguments: {"actor": "Leonardo DiCaprio"}[unused12]"
},
{
    "role": "observation",
    "content": "{"movies": ["Inception", "The Wolf of Wall Street", "Titanic", "The Revenant", "Shutter Island"]}"
},
{
    "role": "reasoner",
    "content": "The function has been successfully executed and five movies starring Leonardo DiCaprio are returned, which can solve the user's query.\n Next: summarizer."
},
{
    "role": "summarizer",
    "content": "Sure, here are some movies starring Leonardo DiCaprio: "Inception", "The Wolf of Wall Street", "Titanic", "The Revenant", "Shutter Island". Enjoy your movies!"
}

All data and code will be open-sourced to contribute to the community. We supplement the experiments of LoRA in two settings: a single LoRA without roles (standard, trained on raw data) and a single LoRA with roles (multi-roles, trained on multi-role data).

Interestingly, the accuracy of LoRA (multi-roles) is lower than that of standard LoRA. The possible reason is that under limited training parameters, there is interference when a single LoRA learns the multi-roles datasets. In contrast, MoR can avoid interference by the architecture of mixture-of-roles to make full use of the limited training parameters and achieve higher accuracy.

3. Confused by equation(6).

Sorry for misleading. The size of u here can be represented by [sequence_length,hidden_dim] (for simplicity, we can omit the dimension of batch size). Without loss of generality, take the size of u in one layer equal to [4096,hidden_dim] as an example. Further example, [:1024, :] of u is processed by reasoner, [1024:1536, :] of u is processed by executor, [1536:2560, :] of u is processed by reasoner, [2560:3072, :] of u is processed by executor, [3072:3584, :] of u is processed by reasoner, and [3584:, :] of u is processed by summarizer. We will further clarify this in the revised version.

Dear Reviewer QZb4:

Sincerely thank you for your constructive review, the concerns and questions are answered in detail.

1. What's the gain by SFT on the same datasets, w/o lora? and other PEFT methods?

Thanks for the helpful suggestion. Based on the same multi-roles dataset, we supplement the experimental results of model Llama3.2-1B-Instruct on BFCL leaderboard. The results are as follows.

From the results, SFT exhibits superior accuracy compared to PEFT methods (LoRA and DoRA), which can be attributed to its more trainable parameters, achieving an average accuracy 10.5% higher than LoRA. Notably, DoRA introduces an advanced scheme by decomposing pretrained weight matrices into magnitude vectors (m) and directional matrices (V), where LoRA is applied specifically to V while m is trained separately. This architectural innovation allows DoRA to surpass LoRA slightly in accuracy. Crucially, our proposed method achieves statistically significant performance improvements through two key innovations: 1) a more rational capacity decomposition strategy, and 2) a novel Mixture-of-Roles framework enabling dynamic interaction between decomposed modules. These enhancements collectively contribute to our method's marked accuracy superiority over SFT and PEFT methods.

2. Validating on more agent structures, such as Reflection, AutoGen.

Thanks for the valuable suggestion. Our method is not constrained by agent architectures. For example, the ​framework of Reflection[1] comprises three modules: ​Actor, ​Self-Reflection, and ​Evaluator. Similarly, in ​AutoGen's[2] application scenarios (Figure 3), multi-agent coding involves modules like ​Commander, ​Writer, and ​Safeguard. While each module in these frameworks operates as an independent LLM, our approach integrates multiple independent LLMs into a ​parameter-efficient Mixture-of-Roles architecture.

Due to the limited time during rebuttal, directly extending our method to these application scenarios proved challenging. Therefore, we modified the agent architecture by adding or removing specific roles based on the existing ​BFCL multi-roles training dataset. The results are as follows.

Our method integrates three core roles: Reasoner, Executor, and Summarizer. Through architectural modifications — specifically removing the Summarizer and introducing a Planner — we extend the framework to 2-roles and 4-roles. Experimental results demonstrate accuracy improvements across all variants compared to the baseline model. Notably, the 3-roles architecture achieves peak performance (accuracy improvement of 50.1% over baseline), whereas the 4-roles, despite containing more trainable parameters, shows diminishing returns (45.7% accuracy gain). This empirical evidence highlights that ​rational role definition and architecture design outweigh mere parameter quantity expansion in achieving optimal model performance.

[1] Reflexion: Language Agents with Verbal Reinforcement Learning

[2] AutoGen: Enabling Next-Gen LLM Applications via Multi-Agent Conversation

3. Validating on more agent-related datasets. How about generalization ability on other agent tasks, maybe refer to AgentGym.

Thanks. In our experiments, we conducted evaluations across multiple datasets including StableToolBench, BFCL, GSM8K, and MATH. In fact, the application scenarios of these datasets involve API calls, Java programming, Python programming, JavaScript programming, and mathematical problem-solving, etc., not just tool use. AgentGym is an impressive work, however, limited by the time of rebuttal, it is difficult for us to reproduce AgentGym in a short time and migrate our methods to it. However, we supplement the experiments of Llama3.2-1B-Instruct on BFCL with varying amounts of training data.

From the results, even with only 1k training samples, we still achieves a 27.4% improvement in average accuracy, demonstrating its strong generalization capability. As the training data volume increases, the accuracy further improves accordingly.