Source-linked AI summary
Chain of Agents: Large Language Models Collaborating on Long-Context Tasks
Yusen Zhang, Ruoxi Sun, Yanfei Chen, Tomas Pfister, Rui Zhang, Sercan Ö. Arik
TL;DR
LLMs struggle with long contexts because retrieval can miss needed information and expanded windows can impair focus. Chain-of-Agents uses sequentially communicating workers over text chunks and a manager synthesizer, and it improves over strong baselines across evaluated long-context tasks by up to 10%.
Problem
LLMs struggle with long-context tasks because input reduction may omit needed information and window extension may weaken focus on pertinent information.
Method
Chain-of-Agents sequentially passes evidence among worker agents processing text chunks, then uses a manager agent to generate the final response.
Results
CoA significantly improves over RAG, Full-Context, and multi-agent baselines across nine datasets covering question answering, summarization, and code completion, by up to 10%.
Takeaways & Limitations
CoA processes the entire input through interleaved reading and reasoning while assigning each agent a short context, mitigating long-context focusing issues.
Takeaways & Limitations
CoA's communication effectiveness, communication forms, and cost or latency remain areas for improvement.
Abstract
from arXiv · showhide
Addressing the challenge of effectively processing long contexts has become a critical issue for Large Language Models (LLMs). Two common strategies have emerged: 1) reducing the input length, such as retrieving relevant chunks by Retrieval-Augmented Generation (RAG), and 2) expanding the context window limit of LLMs. However, both strategies have drawbacks: input reduction has no guarantee of covering the part with needed information, while window extension struggles with focusing on the pertinent information for solving the task. To mitigate these limitations, we propose Chain-of-Agents (CoA), a novel framework that harnesses multi-agent collaboration through natural language to enable information aggregation and context reasoning across various LLMs over long-context tasks. CoA consists of multiple worker agents who sequentially communicate to handle different segmented portions of the text, followed by a manager agent who synthesizes these contributions into a coherent final output. CoA processes the entire input by interleaving reading and reasoning, and it mitigates long context focus issues by assigning each agent a short context. We perform comprehensive evaluation of CoA on a wide range of long-context tasks in question answering, summarization, and code completion, demonstrating significant improvements by up to 10% over strong baselines of RAG, Full-Context, and multi-agent LLMs.
1 Introduction
Long-context tasks challenge LLMs because input reduction can omit needed information while larger windows can hinder focus. Chain-of-Agents addresses both issues by coordinating workers over text segments and having a manager synthesize their evidence, improving performance across evaluated tasks.
- Long-context applications such as question answering, summarization, and code completion can contain entire books or long articles.
- Input reduction, including RAG, may provide incomplete context, while window extension can leave LLMs struggling to focus on pertinent information.
- CoA uses sequential worker communication over segmented text, followed by a manager that synthesizes the accumulated evidence into a final response.
- CoA interleaves reading and processing, assigns each agent a short context, and reduces Full-Context time complexity from n^2 to nk.
- Up to 10%, CoA improves over RAG, Full-Context, and two multi-agent baselines across nine datasets spanning question answering, summarization, and code completion.
2 Related work
Prior multi-agent LLM research emphasizes social simulation and short-text problem solving, while long-context problem solving remains less explored. Existing long-context multi-agent work is closest in using tree structures for single-hop question answering, but lacks broader reasoning and communication capabilities.
- Multi-agent LLM research covers social simulation, games, economies, recommendations, pandemics, and several short-text reasoning tasks.
- The closest cited multi-agent approach uses a tree structure for single-hop question answering over long contexts but lacks multi-hop reasoning and communication between sibling agents.
- Long-context modeling commonly follows input reduction through retrieval or compression, or window extension through model adaptation and larger context limits.
- Complex reasoning research decomposes questions into sub-questions and solves them step by step using prompting, modules, or trained models.
3 Method
CoA divides long inputs into short chunks processed sequentially by communicating worker agents, then uses a manager agent to synthesize the accumulated knowledge into the final answer.
- Stage 1: Worker Agent: CoA splits the source input into l chunks, each short enough to fit within the agent context window.The formulation assumes an input of n tokens and a context limit k, typically k ≪ n.
- Stage 1: Worker Agent: Each worker combines its assigned chunk, the query, worker instructions, and the previous communication unit to produce the next communication unit.This creates a sequential chain in which each worker passes processed information to the next worker.
- Stage 1: Worker Agent: Worker communication units contain task-specific information, including evidence for question answering, summaries for summarization, or code summaries for code completion.This task-dependent content supports CoA across multiple long-context task types.
- Stage 1: Worker Agent: The final worker’s accumulated communication unit spans the full receptive field, allowing CoA to handle varying input lengths by changing the number of workers.Sequential communication lets the last worker access information propagated across the entire input.
- Stage 2: Manager Agent: A manager agent consumes the last worker’s accumulated knowledge and generates the final response, separating chunk analysis from answer production.The paper reports that this separation lets each agent focus on its assigned duty.
- Efficiency: CoA has lower encoding-time complexity than Full-Context when k ≪ n, while both approaches have the same decoding time.The comparison is made theoretically for decoder-only models with average response length r.
4 Experiment
Experiments evaluate CoA across nine long-context datasets, three task types, multiple backbone LLMs, and comparisons with retrieval, full-context, and multi-agent baselines. CoA consistently outperforms these baselines, including long-context models and parallel multi-agent designs.
- Experiment Setup: Experiments cover nine long-context datasets spanning question answering, summarization, and code completion.The evaluation uses six backbone LLMs, including PaLM 2, Gemini 1.0, and Claude 3.
- Experiment Setup: CoA is compared with Vanilla full-context inference, RAG, and two parallel multi-agent approaches: Merge and Hierarchical.RAG reranks 300-word chunks and feeds top-n chunks to the downstream LLM; the multi-agent baselines assign chunks independently to workers.
- Overall Results: CoA outperforms Vanilla and RAG on all eight question-answering datasets across text-bison, text-unicorn, and gemini-ultra.Reported margins over Vanilla (8k) include 13.30% on NarrativeQA, 12.82% on MuSiQue, and 22.00% on QuALITY for the corresponding models.
- Overall Results: CoA outperforms Vanilla at both 8k and 32k context limits on all three summarization and code-completion datasets.For GovReport, RAG with a pseudo-query fails to improve the baseline, whereas CoA improves performance significantly.
- Overall Results: CoA (8k) significantly outperforms Vanilla (200k) on NarrativeQA and BookSum, including samples that fit within the 200k-token window.Improvements over Vanilla (200k) and RAG (8k) increase as Claude 3 scales from Haiku to Opus.
- Overall Results: CoA outperforms Merge and Hierarchical on all eight datasets, while their parallel workers cannot communicate across chunks.Hierarchical beats Vanilla on five of eight datasets, whereas Merge is below Vanilla except on GovReport.
5 Analyses
Analyses show that CoA’s advantages increase when retrieval fails or inputs grow longer, while its sequential design and manager materially affect performance.
- 5.1 CoA Improvement is More Obvious When RAG Fails to Retrieve Gold Answer: When RAG fails to retrieve the gold answer, CoA greatly improves performance beyond RAG’s trend.RAG performs better when the gold answer is nearer the top of its reranked input, but CoA improves substantially when retrieval fails.
- 5.2 CoA Improvement is More Obvious When Long Context Models Meet Longer Inputs: CoA’s improvement over Vanilla becomes more significant as BookSum inputs lengthen, reaching around 100% above 400k tokens.CoA continues improving with longer samples, even against Claude 3’s 200k-token Vanilla baseline.
- 5.4 CoA Enables Complex Reasoning Over Long Contexts: Sequential worker communication enables CoA to perform complex multi-hop reasoning over long contexts in the HotpotQA case study.The first workers explore and broaden related topics before later inference, unlike retrieval based only on query similarity.
- 5.3 CoA Mitigates “Lost-in-the-Middle” Phenomenon: CoA narrows the lost-in-the-middle performance gap from 6.13 (±2.17) for Vanilla to 4.89 (±1.91).The analysis attributes this resilience to giving each agent a shorter context to focus on.
- 5.5 Ablation Study: Effectiveness of Manager and Alternative Design Choices: Multi-path CoA further enhances performance through judge or voting-based augmentation.The comparison evaluates three multi-path augmentation strategies.
- 5.5 Ablation Study: Effectiveness of Manager and Alternative Design Choices: Removing the manager drops MuSiQue performance by more than 10%, while left-to-right reading scores highest on most datasets.The ablation compares direct last-worker generation with alternative reading orders, including Right-to-Left and Permutation.
6 Conclusion
The paper concludes that Chain-of-Agents is a simple, training-free framework for long-context tasks that combines information aggregation with context reasoning. It reports stronger performance than RAG and long-context LLMs, while identifying communication, communication-form, and efficiency limitations.
- Chain-of-Agents is presented as a training-free, task- and length-agnostic, interpretable, cost-effective framework for long-context tasks.
- CoA outperforms RAG and long-context LLMs by a large margin, mitigates lost-in-the-middle, and gains more on longer samples.
- Current LLM alignment is not optimal for inter-agent communication, and CoA does not explore debating or complex discussion.
- CoA’s cost and latency could be reduced through approaches such as replacing some LLMs with more effective models via model routing.
A Proof of Time Complexity
The time-complexity analysis compares encoding and decoding operations for Full-Context LLMs with Chain-of-Agents when processing a source of n tokens.
- For Full-Context LLMs, encoding processes the entire n-token source before response generation.The analysis defines TEnc from the cumulative attention operations over the input positions.
- Full-Context decoding begins after all input is encoded, with the first decoded token attending to n positions.The total decoding operation is expressed for responses whose length r is much smaller than n.
- CoA splits the source into ⌈n/k⌉ chunks, limiting each worker’s encoded context to k tokens.The encoding expression multiplies per-chunk operations by the number of chunks.
- CoA decoding starts after k input tokens and repeats the per-chunk decoding cost across ⌈n/k⌉ chunks.The analysis assumes response length r is much smaller than k for this expression.
B Implementation Details
The experiments use six models through Vertex Model Garden with fixed generation and temperature settings, while following LongBench and SCROLLS task requirements.
- Experiments use Vertex Model Garden 4 API for six models, with temperature 0 except for Self-consistency.Maximum generation is 2048 tokens for Gemini Ultra and 1024 for the other models; RAG reranking runs on A100 GPUs.
- Task-specific requirements follow the original LongBench and SCROLLS setups across nine datasets.
C Example Outputs
The paper illustrates Chain-of-Agents predictions on HotpotQA, QMSum, and RepoBench-P by showing the question, agent communication units, and manager's final prediction.
- Examples from HotpotQA, QMSum, and RepoBench-P present each task's question, agent communication units, and final manager prediction.
D Robustness against Context Window Size
CoA remains effective across agent context-window sizes: performance rises from 4k to 16k and then stabilizes near 20 as the window grows further.
- Performance increases from 4k to 16k context windows on Claude 3 Haiku for NarrativeQA.
- Performance stabilizes at around 20 when the context window increases beyond 16k.
- The evaluation varies Claude 3 Haiku's context window across 4k, 8k, 16k, 32k, 64k, and 128k settings.
E Broader Impacts
CoA is presented as a generic and interpretable framework for complex long-context tasks, but its deployment may require careful prompt design and incur additional API latency.
- CoA is described as a generic framework applicable to diverse long-context tasks beyond those evaluated in the paper.
- Its interpretability lets users check result correctness and may reduce LLM misuse and faults.
- Careful prompt design is required for unseen LLMs, and users may not obtain optimal solutions on newly proposed models.
- Additional API calls can increase network traffic and latency for user pools.