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

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, divideand-conquer frameworks (DCF) split documents into small chunks for independent reasoning and aggregation. While effective for local reasoning, DCF struggles to capture longrange 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-andconquer 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 . Large Language Models (LLMs) with limited context windows (e.g., 8k, 32k) struggle with reasoning over long contexts. As context length increases, the performance declines due to difficulties in processing information far from the text's beginning * Equal contribution. † Corresponding Author. Retrieved Context c 1 c 2 c 3 c i-1 c i c n Chunk c 3 c i c 2 + Reasoning Answer (b) Divide-and-Conquer Framework (a) Retrieval-Augmented Generation Encoder Q Query (c) Our ToM Framework I. Treat Long Context as DocTree II. Perform MapReduce on DocTree Retrieve Query Node 1 Node 2