Stop Using Fixed Chunk Sizes for RAG

AI21 Labs conducted experiments using datasets like QMSum, NarrativeQA, and Seinfeld to test the effects of different chunk sizes. Each dataset was duplicated six times, with each duplication assigned a fixed chunk size ranging from 200 to 1,000 tokens. This methodology allowed for a systematic analysis of how the nature of a query dictates the optimal chunk size for retrieval.

The experiments revealed that the specificity of a query directly influences which chunk size performs best for retrieval. Small chunk sizes, such as 100 tokens, are more effective for specific, contained questions, while broad or diffuse questions, like character descriptions, require larger contexts. This phenomenon was clearly demonstrated using the Seinfeld dataset, where different query types yielded varied results based on chunk size.

To understand the upper limit of retrieval performance, AI21 Labs simulated an 'Oracle' capable of selecting the best chunk size for each individual query. This 'Oracle Experiment' established a benchmark for retrieval, comparing the recall at K against fixed-size chunking. The results indicated a significant performance potential gap ranging from 20% to 40% between fixed-size methods and an ideal adaptive system.

Analysis of multiple datasets revealed that no single fixed chunk size consistently dominates retrieval performance. The intersection of performance lines for various fixed sizes across datasets illustrates this point. This finding underscores that the arbitrary choice of a fixed chunk size significantly impacts retrieval results, highlighting the need for more adaptive strategies.

Indexing, where chunk sizes are determined, occurs without knowledge of future queries. Conversely, during retrieval, queries are known, but the ability to adjust chunk sizes is lost. This disconnect prevents the implementation of adaptive chunking strategies that could optimize retrieval based on query characteristics.

This approach involves committing to several chunk sizes rather than just one. It essentially duplicates and chunks the dataset using multiple window sizes, requiring querying multiple databases at retrieval time. This method contrasts with prior research that focused on enriching the latent spaces of fixed chunks.

A significant hurdle in multiscale indexing is the inherent incomparability of chunks with different sizes within a single ranking system. The speaker highlighted that having 'n' rankings, each based on arbitrarily sized chunks, complicates the process of accurately capturing document relevance. To resolve this, a merging algorithm is proposed to address the issue of varied document representation, enabling effective comparison across different chunk scales.

Instead of retrieving specific chunks, the proposed system retrieves the entire document when a chunk is identified. This paradigm shift creates 'n' rankings of the same documents, making them directly comparable across different chunking scales. By utilizing context windows, the approach provides more comprehensive information, enhancing the overall relevance of retrieved data.

Reciprocal Rank Fusion (RRF) is employed to combine the rankings derived from different chunk-size strategies. This method treats retrieval as a voting mechanism, rather than a singular ranking process, effectively aggregating diverse perspectives. RRF is a simple formula that operates as a lightweight script, adding minimal computational overhead to the retrieval process.

The Multiscale Indexing method demonstrated consistent and significant performance improvements across a variety of datasets, including QM, SAM, Narrative QA, Seinfeld, and Finance Bench. Performance gains ranged from 10% to 40% across various metrics, consistently outperforming traditional fixed-size chunking strategies. Data visualizations, such as greener heat maps, further confirm the superior effectiveness of the multi-chunk approach.

While offering significant performance gains, the Multiscale Indexing approach incurs a cost in terms of memory usage. It requires keeping multiple copies of the database in memory, leading to an estimated memory increase of 2x to 5x. However, the impact on latency is minimal, as retrieval and Reciprocal Rank Fusion (RRF) calculations are highly efficient, maintaining quick response times.

The current selection of chunk sizes, such as 50, 100, and 200 tokens, was acknowledged as largely arbitrary. Future research will focus on optimizing the number of copies needed and determining the most effective chunk sizes through more systematic methods. The team is also exploring alternatives beyond standard Reciprocal Rank Fusion (RRF) to further refine the architectural evolution of the system.

The speaker emphasized that despite advancements in AI, agents have not eliminated the need for effective retrieval systems. The prevailing Retrieval-Augmented Generation (RAG) infrastructure, often based on 2022 standards, is outdated. Simple algorithmic improvements can yield substantial performance gains of 20% to 40% without requiring sophisticated AI models, indicating that RAG infrastructure is ripe for optimization.