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

TL;DR

ToM introduces a hierarchical Tree-oriented MapReduce framework that enhances long-context reasoning in large language models by leveraging document structure for recursive, coherent reasoning over extended texts.

Contribution

It presents a novel hierarchical framework that improves long-context reasoning by constructing a DocTree and performing recursive reasoning with MapReduce, outperforming existing methods.

Findings

Significantly outperforms existing frameworks in long-context reasoning.

Achieves better logical coherence in reasoning tasks.

Effective on 70B+ parameter language models.

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 .