TBDFiltering: Sample-Efficient Tree-Based Data Filtering

TL;DR

This paper introduces a query-efficient, hierarchical clustering method for selecting high-quality training data for large language models, reducing the need for extensive LLM queries and improving data filtering accuracy.

Contribution

It presents a novel text-embedding-based hierarchical clustering approach that adaptively estimates cluster quality with minimal LLM queries, backed by theoretical guarantees.

Findings

Method is query efficient, requiring few document evaluations.

The approach accurately predicts document quality with high probability.

Experimental results show improved filtering over existing classifier-based methods.

Abstract

The quality of machine learning models depends heavily on their training data. Selecting high-quality, diverse training sets for large language models (LLMs) is a difficult task, due to the lack of cheap and reliable quality metrics. While querying existing LLMs for document quality is common, this is not scalable to the large number (billions) of documents used in training. Instead, practitioners often use classifiers trained on sparse quality signals. In this paper, we propose a text-embedding-based hierarchical clustering approach that adaptively selects the documents to be evaluated by the LLM to estimate cluster quality. We prove that our method is query efficient: under the assumption that the hierarchical clustering contains a subtree such that each leaf cluster in the tree is pure enough (i.e., it mostly contains either only good or only bad documents), with high probability, the method can correctly predict the quality of each document after querying a small number of documents. The number of such documents is proportional to the size of the smallest subtree with (almost) pure leaves, without the algorithm knowing this subtree in advance. Furthermore, in a comprehensive experimental study, we demonstrate the benefits of our algorithm compared to other classifier-based filtering methods.