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

TL;DR

This paper introduces Mix-of-Granularity (MoG), a dynamic method for optimizing knowledge source chunking in Retrieval-Augmented Generation systems, improving information retrieval and downstream task performance.

Contribution

The paper proposes MoG and MoG-Graph, novel methods that adaptively determine the best data granularity for knowledge retrieval, with a new training loss and graph-based extension.

Findings

MoG accurately predicts optimal granularity levels.

MoGG enhances retrieval of distant snippets.

Both methods improve RAG system performance.

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.