Skip to content

Facilitating Long Context Understanding via Supervised Chain-of-Thought Reasoning

🕒 Published (v1): 2025-02-18 18:50 UTC · Source: Arxiv · Venue: EMNLP 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

LLMs with extended context windows plateau at 45–55% on practical long-context benchmarks regardless of window size, because raw length extension is insufficient without explicit multi-step reasoning. This paper proposes PAI (Property-based Agentic Inference), a three-agent pipeline that generates Chain-of-Thought–annotated synthetic financial QA data (LongFinanceQA), then distills that reasoning capability into LLaMA-3.1-8B via supervised fine-tuning. GPT-4o-mini augmented with PAI at inference gains +20.3% on Loong; the fine-tuned LongPAI gains +28.0% on Loong's financial subset.

Problem

Existing long-context synthetic datasets pair complex multi-hop questions with brief final answers, omitting intermediate reasoning. This forces models to learn from (long-context query, short answer) pairs with no intermediate signal, leading to suboptimal training. The root issue is formalized as conflating \(p_\theta(A|X,Q)\) with the fuller joint \(p_\theta(R,A|X,Q)\).

Method

PAI decomposes long-context answering into three chained LLM agents:

  1. Property Extraction Agent — given query \(Q\), extracts a set of structured properties \(\{p_i\}_{i=1}^{N_p}\), each a (metric, subject) pair (e.g., {revenue, Company A 2018 report}), using function-calling to enforce structured outputs.
  2. Property-based Retrieval Agent — converts each property into a sub-query \(q_i\), retrieves 1,024-token chunks from long documents \(X\), and generates intermediate findings (sub-answers).
  3. Summarization Agent — synthesizes all intermediate findings into a coherent conclusion \(A\).

The three agents' outputs are concatenated into a <reasoning>…</reasoning> + conclusion format, producing CoT-augmented answers for LongFinanceQA (46,457 QA pairs over 6,911 bilingual financial reports, filtered to 20K–256K tokens).

Supervised CoT Reasoning: LLaMA-3.1-8B-Instruct is first context-extended from 128K→262K via continued pretraining on 1.6B tokens (Slimpajama), then fine-tuned on LongFinanceQA to maximize: $\(\mathcal{L}(\theta) = \sum_{i=1}^{N_r + N_y} \log p_\theta(R, Y | X, Q)\)$ with non-answer positions masked out. Ring attention + flash-attention2 + zigzag sharding enable full fine-tuning at >250K sequence length on 8×A100.

Key Contributions

  • LongFinanceQA: 46,457 long-context QA pairs with explicit CoT reasoning trajectories, derived from 6,911 bilingual (EN/ZH) financial annual reports; single- and multi-source question splits.
  • PAI framework: a structured three-agent annotator that converts free-form long-context questions into property-grounded, verifiable retrieval chains, yielding interpretable and auditable reasoning traces.
  • LongPAI model: supervised distillation of PAI's multi-step agentic inference into a single LLaMA-3.1-8B forward pass via fine-tuning on LongFinanceQA.

Results

  • GPT-4o-mini w/ PAI vs. GPT-4o-mini (base) on Loong overall: +20.3% average score (49.25→59.46 overall PR: 0.24→?; spotlight: +20.2%; comparison: +15.7%; clustering: +29.2%; chain of reasoning: +10.2%).
  • GPT-4o-mini w/ PAI outperforms Gemini-1.5-pro (55.37 vs. GPT-4o-mini w/ PAI ≈ 70+) on Loong overall despite using a substantially weaker base model.
  • LongPAI vs. LLaMA-3.1-8B-Instruct on Loong Financial subset: +28.0% gain; LongPAI surpasses its teacher model PAI in several sub-tasks.
  • Human evaluation of 200 LongFinanceQA samples (1–5 scale, N=10 annotators): final answers and intermediate reasoning steps both score above 4.0 on average, confirming annotation quality.

Limitations

  • LongFinanceQA is domain-specific (financial annual reports); generalization to other long-context domains is not demonstrated.
  • PAI relies on human-crafted agent designs; the three-step workflow is fixed and cannot dynamically adapt to novel query types outside the property (metric, subject) schema.
  • Questions are capped at 256K tokens due to model limits, excluding the longest real-world documents.
  • Context extension pretraining uses Slimpajama (general text), which may not align well with financial domain distributions.
  • PAI evaluation is limited to two benchmarks (Loong, ∞Bench); broader evaluation on other long-context tasks (e.g., SCROLLS, LOFT) is absent.

Relevance to Harnesses / Meta-Harnesses

PAI is a concrete instance of a data-generation meta-harness: a multi-agent orchestration pipeline used not for direct task execution but to fabricate high-quality supervised training signal at scale (46K samples). The three-agent pipeline—extract, retrieve, summarize—exemplifies a fixed-topology harness where each agent's output is the next agent's input, analogous to a preprocessing pipeline run autonomously at dataset-construction time. The paper also demonstrates harness distillation: the expensive agentic harness (multi-LLM, multi-step) is compressed into a single model (LongPAI) via supervised fine-tuning, raising the broader question of when a harness's inference-time cost can be amortized into model weights. This is highly relevant to researchers designing meta-harnesses for labeling or annotation, as it validates the regime where harness-generated data quality exceeds what a direct fine-tuned model can achieve alone.