SmartChunk Retrieval: Query-Aware Chunk Compression with Planning for Efficient Document RAG

TL;DR

SmartChunk retrieval introduces a query-adaptive, reinforcement learning-based framework for dynamic document chunking, improving retrieval accuracy and efficiency in long-document question answering tasks.

Contribution

It proposes a novel planner with reinforcement learning to adaptively select chunk abstraction levels, enhancing retrieval quality and scalability over fixed strategies.

Findings

Outperforms state-of-the-art RAG baselines on multiple QA benchmarks.

Reduces retrieval cost while maintaining high accuracy.

Demonstrates strong scalability with larger corpora and diverse datasets.

Abstract

Retrieval-augmented generation (RAG) has strong potential for producing accurate and factual outputs by combining language models (LMs) with evidence retrieved from large text corpora. However, current pipelines are limited by static chunking and flat retrieval: documents are split into short, predetermined, fixed-size chunks, embeddings are retrieved uniformly, and generation relies on whatever chunks are returned. This design brings challenges, as retrieval quality is highly sensitive to chunk size, often introduces noise from irrelevant or misleading chunks, and scales poorly to large corpora. We present SmartChunk retrieval, a query-adaptive framework for efficient and robust long-document question answering (QA). SmartChunk uses (i) a planner that predicts the optimal chunk abstraction level for each query, and (ii) a lightweight compression module that produces high-level chunk embeddings without repeated summarization. By adapting retrieval granularity on the fly, SmartChunk balances accuracy with efficiency and avoids the drawbacks of fixed strategies. Notably, our planner can reason about chunk abstractions through a novel reinforcement learning scheme, STITCH, which boosts accuracy and generalization. To reflect real-world applications, where users face diverse document types and query styles, we evaluate SmartChunk on five QA benchmarks plus one out-of-domain dataset. Across these evaluations, SmartChunk outperforms state-of-the-art RAG baselines, while reducing cost. Further analysis demonstrates strong scalability with larger corpora and consistent gains on out-of-domain datasets, highlighting its effectiveness as a general framework for adaptive retrieval.