AMOR: A Recipe for Building Adaptable Modular Knowledge Agents Through Process Feedback

Regarding the complexity and human feedback dependency.

We would like to address your concerns as follows:

• Regarding the “Separately Fine-tuned LLMs”: We would like to emphasize that this argument might be inaccurate since we use the same MA-MoE model for all modules and activate different experts to execute their corresponding modules. This approach eliminates the need for separately fine-tuning multiple LLMs.
• Regarding the Two-Stage Fine-Tuning: We argue that the two-stage fine-tuning process is analogous to the widely adopted "pretrain and fine-tune" paradigm. The warm-up stage, which is performed only once, enables AMOR to generalize across different knowledge environments. Once the warm-up stage is complete, AMOR can be deployed and adapted to various domains through the efficient adaptation stage. We believe that the benefits of improved generalization and domain adaptation justify the additional computational cost.
• Regarding the Dependency on Human Feedback: We underscore the importance of maintaining an adaptive mechanism that leverages human feedback to continuously improve performance. Previous agent frameworks often overlook this crucial aspect. AMOR's process feedback mechanism enables efficient and targeted feedback, reducing the overall burden on human supervisors.

In summary, while AMOR introduces complexities, its modular design, two-stage fine-tuning, and process feedback mechanism are deliberate choices that enable it to handle complex tasks effectively and adapt to specific domains efficiently. We believe that the benefits of improved performance and domain adaptability outweigh the concerns raised regarding complexity and human feedback dependency.

Weakness: Regarding the measure of uncertainty.

We agree that it is crucial to provide measures of uncertainty to assess the statistical significance and robustness of the results. To address this concern, we show the mean and standard deviation across three independent runs in the table below.

As shown in the table, the standard deviations across multiple runs are very small, suggesting that our approach yields consistent and robust results. This strengthens the validity of our conclusions and provide confidence in the reported performance.

We will incorporate the uncertainty measures in the revised manuscript to enhance the transparency and credibility of our work.

Question: Regarding module choosing

Each module corresponds to a distinct expert in the MA-MoE model. When AMOR executes a certain module, its module ID will be provided to the routers of the MA-MoE model to indicate which expert should be activated, thereby enabling our model to be “module-aware.” We will endeavor to provide a more comprehensive explanation of this module-expert mapping mechanism in the revised version of our paper.

Weakness 1: Regarding the scope and problem formulation

• Scope. AMOR aims to develop a general framework for building adaptable modular LLM agents that can leverage external knowledge sources to tackle complex reasoning tasks. However, we appreciate the reviewer's advice that being more explicit in framing the specific tasks and domains upfront would strengthen our readability. We will clarify the scope early in the next revision.
• Problem Formulation. We also appreciate the reviewer's advice to provide a clear task formulation upfront. In the revision, we will add a problem statement and formulation section, which will precisely define the QA-style reasoning tasks we use AMOR to solve.

Weakness 2: Regarding the method description

In Equations 1 and 3, the policy $\pi$ refers to the strategy of the MA-MoE model that maps from the state $s$ to an action $y$.

We will explicitly define it to improve clarity in the revision.

Weakness 3: Regarding simpler presentation

AMOR uses several novel techniques, including FSM-based reasoning logic, process feedback mechanism, and two-stage fine-tuning strategy. The paper is organized to provide clear definitions, examples, and explanations without oversimplifying the technical details. We will simplify the content in the revision and welcome specific advice from the reviewer on how we could improve the clarity.

Weakness 4: Regarding the details of AMOR overcoming challenges in the introduction

We believe we have made every effort to convey the key insights and intuitions behind AMOR and how it addresses the stated challenges.

As summarized in Tab. 1, current agents face challenges in three main aspects: (1): uncontrollable reasoning logic; (2): lack of adaptation mechanisms for new environments; and (3): difficulty for humans to intervene in the reasoning process.

In the fourth paragraph (Lines 34-41), we outline the core idea behind AMOR. The FSM-based reasoning logic enables AMOR to solve problems via executions and transitions over a set of modules, allowing for process-based feedback. This design directly addresses challenges (1) and (3).

In the fifth paragraph (Lines 42-50), we provide technical details on AMOR's adaptive mechanism, which tackles challenge (2).

We welcome any further feedback from the reviewer on how we could improve the presentation of our work in the revision.

Weakness 5: Regarding process-based feedback

Alg. 1 shows that AMOR solves problems through executions and transitions over a set of modules. In each reasoning step, AMOR executes the module $m$ while in a specific state $s$, obtains the output $y$, and transits to the next state. The overall reasoning process $R$ is formed by a series of such steps. Each step can receive human feedback, termed process-based feedback, as opposed to outcome feedback typically provided for the final step. We will elucidate this concept more explicitly in the revision.

Weakness 6: Regarding the experimental setup

We provide additional details about the experimental setup below:

• Without Fine-tuning: We apply different methods directly to off-the-shelf LLMs without fine-tuning. We provide in-context examples to instruct LLMs to solve given problems. Tab. 12-15 in Appendix A.1 show the prompts for four LLM modules in AMOR under this setting on HotpotQA.
• With Fine-tuning: LLMs are fine-tuned on specific datasets. For AMOR, we employ a two-stage fine-tuning strategy:
  • Warm-up Fine-tuning: The stage fine-tunes an LLM on trajectories from public datasets.
  • Adaptation Fine-tuning: The stage allows AMOR to adapt to specific environments by leveraging different forms of feedback:
    • Process Feedback: AMOR is optimized using feedback provided for both intermediate and final steps of its reasoning process.
    • Outcome Feedback: AMOR is optimized using feedback provided only for the final step of its reasoning process.

We will include these clarifications in our revision.

Weakness 7: Regarding the table and figure size

We will explore ways to enhance their clarity in the revision.

Question 1: Regarding the fairness between AMOR and baselines

We have carefully designed our experimental setup to ensure fair comparisons in terms of training data sizes. It involves two groups of comparisons:

• Comparisons between AMOR$_{\rm WFT}$ and baselines that are NOT fine-tuned on the target datasets: The table below lists the training data sources and sizes for AMOR$_{\rm WFT}$ and baselines.

  Methods	Data Source	Size
  FireAct	GPT-4 Trajectories	2.5k
  AgentLM	GPT-4/3.5 Trajectories	1.8k
  Self-RAG	Public Resources	145k
  LUMOS	Public Resources	57k
  AMOR$_{\rm WFT}$	Public Resources	50k

  • While FireAct and AgentLM use fewer data than AMOR$_{\rm WFT}$, they rely on the GPT-4/3.5 API for data annotation, which hinders them from scaling up training examples. We believe that AMOR's ability to achieve sota results without any proprietary LLMs is a substantial milestone.
  • Compared to Self-RAG and LUMOS, which also use public resources, AMOR$_{\rm WFT}$ uses fewer or a comparable number of training examples.

  In summary, this group of comparisons is fair, as AMOR$_{\rm WFT}$ uses fewer or a comparable number of data from public resources without relying on proprietary LLMs.

• Comparisons between AMOR$_{\rm Process}$ and baselines that are fine-tuned on the target datasets: AMOR$\_{\rm Process}$, AMOR$_{\rm Outcome}$ and OneR use the same training data from the target datasets, leading to a fair comparison. Table 3 shows the data statistics.

We will provide detailed explanations in the revision.

Question 3: Regarding the difference between AMOR and prior modular agents

Please kindly refer to "Author Rebuttal by Authors" provided at the very beginning.

Weakness 1: Regarding the difference between AMOR and prior reasoning methods

Please kindly refer to "Author Rebuttal by Authors" provided at the very beginning.

Weakness 2: Regarding the technical details.

We acknowledge the importance of transparency and reproducibility in scientific research, and therefore we have included the implementation details, including hyper-parameters, in Lines 204-210 of our manuscript. Additionally, we have submitted the source code for MA-MoE implementation as part of the supplementary material. We intend to release a publicly accessible code repository after the double-blind review period. This repository will facilitate follow-up research to replicate and extend our work.

To further address the reviewer's comments, we plan to include an additional supplementary section in the final version of our paper, which will provide a comprehensive explanation of the algorithm underlying AMOR, elucidating the reasoning process in detail.

We believe that these actions will adequately address the reviewer's concerns, thereby enhancing the clarity and reproducibility of our model.

Weakness 3: Regarding the training details

Equation 3 is the weighted sum of the standard RLHF optimization objective over all modules. Each module corresponds to a parameterized policy network $\pi_{\theta_m}$. The optimization objective is to maximize the reward with a KL divergence penalty. In this paper, we adopt the KTO algorithm [1] for optimization. Equation 3 illustrates the gradients for KTO.

We brief the derivation from Equation 3 to a differentiable loss as follows: As introduced by prior works [1, 2], it is straightforward to show that the closed-form expression for the optimal reward function $r^*_m$ for each module $m$ takes the form: $r^*_m=\beta\text{log}({\pi\_{\theta_m}}/{\pi\_{\theta_m}^{\text{old}}})+\beta\text{log}Z_m$, where $Z_m$ is the partition function. Applying the expression to the Kahneman-Tversky model [1], we can get the differentiable KTO loss. Please kindly refer to the KTO paper for more details. We will provide detailed descriptions to improve the clarity of my work in the revision.

[1] Ethayarajh et al. KTO: Model Alignment as Prospect Theoretic Optimization. ICML 2024.

[2] Rafailov et al. Direct Preference Optimization: Your Language Model is Secretly a Reward Model. NeurIPS 2023.

Weakness 4: Regarding the gradients from the feedback

In the MA-MoE architecture, different modules correspond to different FFN layers and share the same embedding layers and multi-head self-attention layers. Therefore, regarding a training example from $\mathcal{R}_m$ for module $m$, the gradient induced by the feedback will go through the whole MA-MoE model, except those FFN layers corresponding to other modules.