Technical Summary: AMATA – Adaptive Multi-Agent Trajectory Alignment for Knowledge-Intensive Question Answering

1. Problem Statement

Despite significant advancements in Large Language Models (LLMs), generating factually consistent responses for knowledge-intensive Question Answering (QA) remains challenging. Current models struggle with hallucinations and limitations in bridging long-tail knowledge gaps. While Retrieval-Augmented Generation (RAG) offers a mitigation strategy by accessing external knowledge, it suffers from retrieval inaccuracies and increased inference latency due to longer contexts.

Existing multi-agent approaches for knowledge-intensive QA generally fall into three paradigms, each with distinct limitations:

1. Modular Training: Fine-tunes agents independently on subtasks. This lacks global optimization, leading to error propagation and an inability to compensate for inter-agent dependencies.
2. End-to-End Training: Jointly optimizes all agents within a single framework. While enabling co-adaptation, uniform gradient updates often obscure the distinct specialization demands of heterogeneous agents, causing over-homogenization.
3. Global-Local Training: Combines independent subtask optimization with joint fine-tuning. However, these methods often fail to capture dynamic inter-agent dependencies, which are crucial for diverse tasks (e.g., unnecessary verification steps for high-confidence questions, or rigid execution sequences that ignore context).

The core problem addressed is the lack of a framework that can dynamically align agent participation based on specific question requirements while modeling the complex, context-dependent dependencies between agents to improve both reasoning performance and token efficiency.

2. Methodology: AMATA Framework

The authors propose AMATA (Adaptive Multi-Agent Trajectory Alignment), a framework designed to improve agent-level alignment and capture dynamic inter-agent dependencies. The system utilizes six specialized agents: Intent Reconstructor, Knowledge Retriever, Knowledge Filter, Knowledge Locator, Response Generator, and Answer Verifier.

The methodology is grounded in two primary innovations:

A. Intra-Trajectory Preference Learning

Existing approaches often treat all agents as uniformly relevant. AMATA introduces a mechanism to dynamically optimize agent participation for each specific question.

• Mechanism: The system learns question-specific agent preference distributions. For each agent-question pair, a preference score (e.g., <Retriever: 5> vs. <Verifier: 1>) is assigned based on the agent's utility for the current input.
• Implementation: These scores are concatenated with agent descriptions as a prefix and paired with the question for fine-tuning. This allows a monolithic model to capture heterogeneous agent competencies through Supervised Fine-Tuning (SFT) augmented with adaptive prefix scoring, rather than requiring separate datasets for each agent.
• Goal: To prioritize critical agents (e.g., knowledge tools for complex queries) and suppress redundant ones (e.g., verifiers for high-confidence answers), thereby reducing token consumption.

B. Inter-Agent Dependency Learning

Multi-agent systems exhibit context-dependent dependencies where triggering a pivotal agent necessitates the coordinated execution of linked agents.

• Mechanism: The authors introduce a Dependency-aware Direct Preference Optimization (DA-DPO) module. Unlike standard DPO which compares single winning/losing pairs, DA-DPO learns context-sensitive execution rankings.
• Implementation:
  • Preference samples are constructed to explicitly encode inter-agent dependencies.
  • Trajectories are annotated with a joint preference score reflecting the global optimality of the multi-agent sequence.
  • The algorithm selects top-$K$ winning samples based on "dependency goodness" (how well the agent scores correlate with successful outcomes) and treats the remaining samples as losing examples.
  • A listwise Plackett-Luce preference model is used to optimize the policy, prioritizing trajectories with robust inter-agent coordination.
• Goal: To enable the LLM to infer optimal multi-agent execution sequences with high reliability, ensuring that dependent agents (e.g., Retriever $\rightarrow$ Filter $\rightarrow$ Locator) are executed in a coherent, dependency-aware manner.

Training Pipeline

The framework undergoes a two-stage training process:

1. Intra-Trajectory Learning: Optimizes the model to perceive agent utility via the preference prefix mechanism ($L_{Intra}$).
2. Inter-Agent Dependency Learning: Combines the basic agent prediction loss with the DA-DPO loss ($L_{Inter}$) to enhance multi-agent cooperation ($L_{total} = \alpha_1 L(T) + \alpha_2 L_{Inter}$).

3. Key Contributions

1. Adaptive Agent Customization: A novel approach to dynamically modulate agent influence based on question complexity, moving away from static workflows or uniform agent relevance.
2. Dependency-Aware Optimization: The introduction of DA-DPO, which captures cross-agent tool dependencies and learns relative execution rankings, addressing the "error propagation" and "rigid sequencing" issues in prior multi-agent systems.
3. Efficiency and Performance: A framework that achieves high reasoning performance while significantly reducing token overhead by dynamically skipping unnecessary agents (e.g., verifiers) and optimizing the execution path.

4. Experimental Results

AMATA was evaluated on five established knowledge-intensive QA benchmarks: HealthQA, ARC-Choice, PopQA, SQuAD 1.1, and ASQA.

• Performance: AMATA consistently outperformed baseline approaches, including vanilla QA methods, knowledge-augmented frameworks (e.g., Self-RAG, RADIT), and LLM-based trajectory systems (e.g., SMART, GiGPO, SPA-RL). It achieved an average performance improvement of +4.02% across all tasks.
• Efficiency: The method reduced token consumption overhead by approximately 70% compared to strong baselines. This is attributed to the adaptive preference learning which avoids invoking unnecessary agents.
• Ablation Studies: Removing the inter-agent dependency loss ($L_{Inter}$) resulted in the steepest performance decline, confirming the importance of modeling agent associations. Similarly, removing specific agents (like the Filter or Verifier) degraded performance, validating the necessity of the full agent set and their specific roles.
• Token Analysis: Compared to global-local baselines like SMART, AMATA maintained comparable token overhead but significantly outperformed it in accuracy. RL-based methods (e.g., GiGPO) showed higher token consumption (+70%) due to multiple rollout sessions.

5. Significance and Claims

The paper claims that AMATA represents a significant step forward in multi-agent QA by shifting from static or uniformly optimized agent interactions to adaptive, dependency-aware trajectory alignment.

• Dynamic Adaptation: The framework successfully demonstrates that questions with high confidence do not require verification steps, and complex knowledge questions require specific chains of knowledge agents. AMATA learns these dynamics automatically.
• Scalability and Efficiency: By reducing token consumption by ~70% without sacrificing accuracy, the method addresses the practical cost barriers of deploying complex multi-agent systems.
• Generalization: The authors note that the framework operates as a zero-shot, self-adaptive system at inference, requiring no pre-existing training traces for new domains once the "collaboration policy" is learned.

The authors maintain modesty regarding limitations, acknowledging that experiments were primarily conducted on 7B-parameter models and that the number of winning/losing samples for DPO was limited. They suggest that scaling to larger models (e.g., 70B) and increasing sample diversity could further enhance performance on complex tasks.