Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Mix-of-Granularity: Optimize the Chunking Granularity for Retrieval-Augmented Generation

Zijie Zhong, Hanwen Liu, Xiaoya Cui, Xiaofan Zhang, Zengchang Qin

TL;DR

This work tackles the challenge of fixed chunking in retrieval-augmented generation by introducing Mix-of-Granularity (MoG), which uses a router to dynamically select the optimal chunk granularity for each query. It further extends to MoGG, reformatting references as graphs to retrieve distantly located snippets across sources. A soft-label loss is proposed to enable end-to-end training despite top-$k$ retrieval, and experiments on medical QA benchmarks show MoG and MoGG achieve meaningful accuracy gains, especially for weaker LLMs, with MoGG offering further improvements through graph-context retrieval. The methods balance precision and recall while incurring manageable time and storage overhead, and the work provides code for reproduction and further research in dynamic chunking for RAG.

Abstract

Integrating information from various reference databases is a major challenge for Retrieval-Augmented Generation (RAG) systems because each knowledge source adopts a unique data structure and follows different conventions. Retrieving from multiple knowledge sources with one fixed strategy usually leads to under-exploitation of information. To mitigate this drawback, inspired by Mix-of-Expert, we introduce Mix-of-Granularity (MoG), a method that dynamically determines the optimal granularity of a knowledge source based on input queries using a router. The router is efficiently trained with a newly proposed loss function employing soft labels. We further extend MoG to MoG-Graph (MoGG), where reference documents are pre-processed as graphs, enabling the retrieval of distantly situated snippets. Experiments demonstrate that MoG and MoGG effectively predict optimal granularity levels, significantly enhancing the performance of the RAG system in downstream tasks. The code of both MoG and MoGG are released in https://github.com/ZGChung/Mix-of-Granularity.

Mix-of-Granularity: Optimize the Chunking Granularity for Retrieval-Augmented Generation

TL;DR

This work tackles the challenge of fixed chunking in retrieval-augmented generation by introducing Mix-of-Granularity (MoG), which uses a router to dynamically select the optimal chunk granularity for each query. It further extends to MoGG, reformatting references as graphs to retrieve distantly located snippets across sources. A soft-label loss is proposed to enable end-to-end training despite top- retrieval, and experiments on medical QA benchmarks show MoG and MoGG achieve meaningful accuracy gains, especially for weaker LLMs, with MoGG offering further improvements through graph-context retrieval. The methods balance precision and recall while incurring manageable time and storage overhead, and the work provides code for reproduction and further research in dynamic chunking for RAG.

Abstract

Integrating information from various reference databases is a major challenge for Retrieval-Augmented Generation (RAG) systems because each knowledge source adopts a unique data structure and follows different conventions. Retrieving from multiple knowledge sources with one fixed strategy usually leads to under-exploitation of information. To mitigate this drawback, inspired by Mix-of-Expert, we introduce Mix-of-Granularity (MoG), a method that dynamically determines the optimal granularity of a knowledge source based on input queries using a router. The router is efficiently trained with a newly proposed loss function employing soft labels. We further extend MoG to MoG-Graph (MoGG), where reference documents are pre-processed as graphs, enabling the retrieval of distantly situated snippets. Experiments demonstrate that MoG and MoGG effectively predict optimal granularity levels, significantly enhancing the performance of the RAG system in downstream tasks. The code of both MoG and MoGG are released in https://github.com/ZGChung/Mix-of-Granularity.
Paper Structure (37 sections, 2 equations, 6 figures, 11 tables)

This paper contains 37 sections, 2 equations, 6 figures, 11 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Retrieval-Augmented Generation
  4. Chunking Optimization
  5. Graph-Based Text Processing
  6. Methodology
  7. Preliminaries
  8. Naive MoG
  9. Multi-granularity Router
  10. Soft-labels
  11. MoGG: MoG with Graph-context
  12. Experiments
  13. Corpus and Medical QA datasets
  14. Experiment Setup
  15. Performance of MoG on MQA Task
  16. ...and 22 more sections

Figures (6)

  • Figure 1: MoG automatically selects the optimal granularity when extracting information from the reference database (scenarios 1 and 2), achieving both high pertinence and coverage. When relevant information is dispersed across distant sections (scenario 3), the reference documents are pre-processed as graph, then MoGG is applied to retrieve these separate snippets from the best hopping range.
  • Figure 2: MoG mechanism prioritizes the chunks retrieved from optimal granularity level, which is determined by the router based on the user input query.
  • Figure 3: Pre-processing the reference document to form graphs. The concept of "granularity level" is changed into "hopping range" in graphs.
  • Figure 4: Averaged weights of different granularity levels on different QA datasets
  • Figure 5: Accuracy of Medical Question-Answering task with different number of candidate snippets
  • ...and 1 more figures