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

Summary

This paper proposes AMOR, which is a modular framework for answering questions by reasoning over external knowledge bases. AMOR is designed as a Finite State Machine (FSM), which provides a structured way to break down the question into smaller pieces and solve complex tasks through various steps. Each state in AMOR is a module, which can be an LLM call or a tool call. Each LLM module is fine-tuned separately using data relevant to only that step, thus making data collection and evaluation easier.

Strengths

The authors propose an innovative framework that incorporates FSMs, which is a "classical" technique, with state-of-the-art LLM and tool calling that gracefully incorporates human feedback. The text is well-written and clearly explains all parts of the framework. The figures are also helpful in understanding their proposal. I can see AMOR framework being useful in real-world scenarios like question-answering in private datasets. The paper also presents extensive experiments using popular datasets like HotPotQA and PubMedQA.

Weaknesses

The weaknesses I see are the complexity of this framework: many moving parts, including various LLMs that need to be fine-tuned separately; and the 2 stages fine-tuning approach (warm-up + fine-tuning) makes it even more complex and resource intensive. Dependency on human feedback is another weakness I can see here. However, these are weaknesses common to other agents frameworks so I wouldn't consider this as a disqualifier for this paper.

Questions

I don't have any questions

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors were transparent about possible limitations and negative impacts in the **Broader impacts and safeguards** session and also included suggestions on how to mitigate these problems.

Summary

This work proposes AMOR, a modular approach to building knowledge agents using open-source LLMs. AMOR decomposes tasks into reasoning logic, represented as a finite state machine (FSM) composed of sequentially chained modules. These modules include tools e.g, document retrieval and LLMs, e.g., answer extraction. AMOR retrieves relevant documents for a query and combines information from different sources to produce an answer. The development process of AMOR includes three stages: 1. **Warm-up Phase:** Each module is fine-tuned on datasets containing not only input-output pairs but also the relevant intermediate steps for each module. This phase ensures that AMOR can generalize across various tasks and knowledge environments. 2. **Adaptation Phase:** The agent is further fine-tuned on specific domains using process feedback. This feedback, which can be human-generated or derived from evaluation benchmarks, is provided at each reasoning step to refine the agent's performance. Empirical evaluation demonstrates that AMOR effectively utilizes the rich warm-up and evaluation datasets, outperforming relevant baselines. Ablation studies confirm that each component of AMOR is essential for achieving maximal performance.

Strengths

This work is well-motivated and addresses a need for knowledge agents. The modular FSM-based approach is innovative and sensible, allowing for precise reasoning logic. The paper is well-written, and the empirical study is thorough and well-conducted, demonstrating the method's effectiveness through extensive experiments and ablation studies.

Weaknesses

All experiments do not report any measure of uncertainty. Thus, it is impossible to determine if the conclusions are statistically significant. Including measures of uncertainty would strengthen the validity of the results.

Questions

How are the modules chosen in MA-MoE? Is the module/step ID provided to the routers?

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

Yes

Summary

This work presents a modular pipeline for QA tasks. The pipeline consists of several modules such as question decomposition, document/passage retrieval, answer extraction, etc. Training data is separately constructed for each module (based on existing datasets) and models are individually fine-tuned for the respective modules. This processes is referred to as the 'warm-up' stage. Models are further updated based on feedback obtained during inference using an alignment algorithm (KTO is used in this work). Experiments on QA datasets shows that the proposed method performs better than various baselines in the literature. Several ablations are also provided to study the impact of each component.

Strengths

* The proposed modular approach is sensible and performant. * Extensive baselines were considered. * Several ablations were performed that show the impact of various components. It is interesting that the KTO based approach performed better than a simple SFT approach. The human feedback experiment was also interesting. * Strong performance on three QA benchmarks (HotpotQA, PubmedQA, Qasper). Overall I think this paper presents interesting ideas. Various presentation aspects can be improved (see weaknesses).

Weaknesses

* The scope for this work was not properly introduced. The paper is motivated from very broad goals only for the reader to later find that the approach and experiments focus on QA tasks. In fact, the paper does not present a formal task description or problem formulation before presenting the methods, leaving the reader to piece together things. * Method description: In addition to the task not being formally presented, several notations were not clearly defined. For example, equations 1, 3 start talking about a policy which was not described earlier. Although I do understand what's happening in these equations, clearly defined problem and terminology would have made things much clearer. * In general, I felt the ideas in the paper could've been framed/described in a much simpler way. Readability could be improved. * The introduction provides motivations but fails to convey any details/intuitions of the proposed method or how it overcomes the challenges mentioned. * The 'processed based feedback' was not clear to me. I think the paper tries to squeeze in many ideas but this also distracts the reader from the core ideas. * The experimental setup could have been made clearer, including the various settings such as with/without finetuning and process/outcome feedback. * Figures and tables are just too small and impossible to parse.

Questions

* What is the without FT setting in Table 6? * Are the comparisons against baselines fair? Does the proposed approach use more data compared to the baselines? * What is the major delta of this paper compared to prior modular reasoning works (e.g., Rewoo, Lumos)?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

I did not find a discussion of limitations in the papers, I suggest the authors to add some discussion.

Summary

This paper proposed an architecture for advanced reasoning in LLMs. The architecture contains several modules dedicated to different tasks in the reasoning flow, each of which can be trained separately using related datasets constructed from public datasets. The proposed method disentangles the reasoning process into sub-steps that are easier to train and suffer less from sparse feedback which is a critical problem for reasoning in traditional LLMs. In the experiments part, the authors conducted sufficient empirical studies, showing the effectiveness of the architecture proposed.

Strengths

1. There's a clear analysis of the bottlenecks that limit the reasoning ability of LLMs. The architecture proposed is designed to solve these problems accordingly. The motivation is clear and the method is reasonable. 2. The structure of the proposed method is well-explained with text and pictures, making it easy for reviewers and readers to understand. 3. In experiments, the choice of baselines has good coverage. The dataset selection respects each subsection's theme, and the experiment setup is reasonable. The experiments are well explained, making it easy to reproduce them.

Weaknesses

1. Currently, many works focus on enhancing the reasoning capability of LLMs. Although the authors mentioned previous works about LLM reasoning and RAG in the related works part, an in-depth explanation of why this method surpasses its predecessors is still needed. A good way to do this is to explain clearly, probably with a few pictures or equations, how those methods dealt with the difficulties and their advantages/drawbacks, and then show that AMOR is better as it avoids some of the existing drawbacks in previous methods. 2. The state transition diagram is well displayed. However, when readers look into the implementation of each module, some key information such as the architecture, hyperparameters, and techniques used for better performance is missing, making it more difficult to understand and reproduce the whole model. A table of these details or a later public code repository would be appreciated. 3. In section 3.3, the detail of how the policy is trained is not well-explained. Did the author use parameterized policies such as policy networks and use a policy-based method in RL for optimization? If so, is the loss presented in equation (3) an auxiliary loss to the loss in policy gradient methods or the pure source of gradients? 4. Following the above question. Will the gradient induced by feedback go through the whole model or only certain blocks?

Questions

See "Weaknesses" part. My concerns listed there are also my questions to the authors.

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors discussed the limitations of their method adequately in the "Limitations" part in the appendix.