ToM: Leveraging Tree-oriented MapReduce for Long-Context Reasoning in Large Language Models

TL;DR

This work tackles the challenge of reasoning over ultra-long contexts in large language models by introducing ToM, a tree-oriented MapReduce framework. ToM builds a DocTree by applying a Hierarchical Semantic Parser to document chunks and aggregating subtrees bottom-up, then performs recursive MapReduce reasoning across the hierarchy to generate and reconcile rationales from leaves to the root. The approach addresses limitations of RAG and divide-and-conquer by preserving cross-chunk relationships and enabling conflict-resolving aggregation across siblings and ancestors. Experimental results on 70B+-parameter LLMs show ToM achieving superior long-context reasoning performance over baselines, with notable gains on ultra-long QA and multi-choice tasks, validating the method’s effectiveness and scalability for complex long-document understanding. The work demonstrates that structured, hierarchical reasoning can significantly improve coherence and information integration in long-context scenarios, highlighting a practical pathway toward more robust long-context AI systems.

Abstract

Large Language Models (LLMs), constrained by limited context windows, often face significant performance degradation when reasoning over long contexts. To address this, Retrieval-Augmented Generation (RAG) retrieves and reasons over chunks but frequently sacrifices logical coherence due to its reliance on similarity-based rankings. Similarly, divide-and-conquer frameworks (DCF) split documents into small chunks for independent reasoning and aggregation. While effective for local reasoning, DCF struggles to capture long-range dependencies and risks inducing conflicts by processing chunks in isolation. To overcome these limitations, we propose ToM, a novel Tree-oriented MapReduce framework for long-context reasoning. ToM leverages the inherent hierarchical structure of long documents (e.g., main headings and subheadings) by constructing a DocTree through hierarchical semantic parsing and performing bottom-up aggregation. Using a Tree MapReduce approach, ToM enables recursive reasoning: in the Map step, rationales are generated at child nodes; in the Reduce step, these rationales are aggregated across sibling nodes to resolve conflicts or reach consensus at parent nodes. Experimental results on 70B+ LLMs show that ToM significantly outperforms existing divide-and-conquer frameworks and retrieval-augmented generation methods, achieving better logical coherence and long-context reasoning. Our code is available at https://github.com/gjn12-31/ToM .

Figures (12)

• Figure 1: Comparison between ToM and existing approaches: LLMs enhanced with RAG (a) rely on sequential reasoning over retrieved chunks, while current Divide-and-Conquer frameworks (b) process chunks in isolation. In contrast, ToM (c) leverages the hierarchical structure of DocTree for tree-based reasoning, mitigating conflicts and preserving contextual coherence through recursive MapReduce reasoning.
• Figure 2: Illustration for DocTree Construction: The process begins with preparation for compression, where a retriever selects Top-k chunks as the foundation for the tree. Each chunk is then processed (A) using the Hierarchical Semantic Parser to capture its internal hierarchy. Next, (B) the hierarchical structure of each chunk is parsed into a subtree, with nodes capturing structured information. These subtrees are merged (C) by placing root nodes at the same level. Following this, (D) a Bottom-up Aggregation integrates information across levels. Finally, the complete DocTree is constructed (E) by combining low-level subtrees with higher-level summaries, ready for reasoning.
• Figure 3: Illustration of the order of recursive reasoning. Nodes at the same hierarchy level, like D and E, can be processed in parallel for accelration.
• Figure 4: Effect of the number of selected chunks.
• Figure 5: Effect of the chunk size.