Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

TBDFiltering: Sample-Efficient Tree-Based Data Filtering

Robert Istvan Busa-Fekete, Julian Zimmert, Anne Xiangyi Zheng, Claudio Gentile, Andras Gyorgy

TL;DR

TBDFiltering tackles the scalability challenge of quality-filtering training data for large language models by combining text embeddings with hierarchical clustering to guide adaptive LLM prompting. The core method greedily traverses a clustering tree, querying only a small, strategically chosen subset of leaves to determine keep/discard decisions, with theoretical guarantees that the required samples scale with the final tree cut complexity $K$ rather than the dataset size. Empirically, TBDFiltering yields consistent downstream improvements (1–5% relative) across ThePile, C4, and FineWeb datasets and model sizes from 270M to 4B, often outperforming baseline and classifier-based filtering approaches. The approach promises scalable, modality-agnostic data curation for high-quality foundation models and invites future work on multi-modal and diversity-aware extensions.

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.

TBDFiltering: Sample-Efficient Tree-Based Data Filtering

TL;DR

TBDFiltering tackles the scalability challenge of quality-filtering training data for large language models by combining text embeddings with hierarchical clustering to guide adaptive LLM prompting. The core method greedily traverses a clustering tree, querying only a small, strategically chosen subset of leaves to determine keep/discard decisions, with theoretical guarantees that the required samples scale with the final tree cut complexity rather than the dataset size. Empirically, TBDFiltering yields consistent downstream improvements (1–5% relative) across ThePile, C4, and FineWeb datasets and model sizes from 270M to 4B, often outperforming baseline and classifier-based filtering approaches. The approach promises scalable, modality-agnostic data curation for high-quality foundation models and invites future work on multi-modal and diversity-aware extensions.

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.
Paper Structure (16 sections, 2 theorems, 8 equations, 4 figures, 3 tables, 1 algorithm)

This paper contains 16 sections, 2 theorems, 8 equations, 4 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Problem Formulation
  3. Algorithm and Theoretical Analysis
  4. Implementation Details
  5. Text representation
  6. Affinity clustering
  7. Ablation
  8. Datasets and statistics
  9. Prompting
  10. Experimental Results
  11. Relative improvement of model performance
  12. Where does the improvement come from?
  13. Why does clustering help?
  14. Conclusion and Ongoing Work
  15. Prompt Template
  16. ...and 1 more sections

Key Result

Proposition 3.1

alg: greedy will produce a filtering $\mathcal{X}_{\text{disc}},\mathcal{X}_{\text{keep}}$ such that for any $\delta \in (0,1)$, with probability at least $1-\delta$, and where $K$ is a random variable denoting the complexity of the final tree cut. The sample complexity of the algorithm is bounded by $2K N_{\max}$.

Figures (4)

  • Figure 1: A hierarchical clustering of the data items $\mathcal{X} = \{x_1,\ldots, x_8\}$. A cut of the tree is depicted as the dotted line. This is a collection of subtrees whose leaves partition $\mathcal{X}$. In this example, the partition induced by the depicted cut is made up of the four subtrees rooted at the blue nodes $n_1,\ldots,n_4$ corresponding to the clusters $D(n_1) = \{x_1,x_2,x_3\}$, $D(n_2) = \{x_4\}$, $D(n_3) = \{x_5,x_6\}$, and $D(n_4)=\{x_7,x_8\}$. The value of the 0/1 quality label $f(\cdot)$ is also depicted inside the leaves. Assumption \ref{['ass:1']} is satisfied with $\alpha' \geq 1/3$, and $\beta'=0$, as the subtrees have average $f$ value of $\bar{f}(n_1) = 1/3$, $\bar{f}(n_2) = 0$, $\bar{f}(n_3) = 0$, and $\bar{f}(n_4) = 1$.
  • Figure 2: Prompt feedback distribution computed based on 100k randomly selected chunks from the three datasets and two prompts we used with TBDFiltering. We indicated the failure of the prompt response by $-1$ which happened very rarely. In our implementation, we assigned to these cases 0 score. The two red crosses in the two histograms correspond to the values of thresholds $\alpha$ and $1-\beta$ that have been selected in the two cases, multiplied by 5.
  • Figure 3: The heatmap of average relative improvement comparing different training setups, specifically when no filtering is applied, and when filtering is carried out through CB(FW-EDU), TBDF(FW-EDU) and TBDF(General). The results for ThePile, C4 and FineWeb datasets are shown in Figure \ref{['fig:thepile']}, \ref{['fig:c4']} and \ref{['fig:fineweb']}, respectively. In each subfigure, the performance of the Gemma 3 model with parameters 270M, 1B and 4B is shown from left to right. Each cell displays the relative improvement score averaged over the 8 evaluation tasks mentioned in Section \ref{['sec:ablation']}. Numbers within brackets indicate the count of evaluation tasks where the model outperformed versus underperformed against the baseline. For example, a result of 6/2 signifies that the model surpassed the baseline on six tasks while performing inferiorly on two. The cells are colored depending on the relative difference in a pairwise comparison: green means that the ablation is better with the corresponding filtering, red indicates inferior performance.
  • Figure 4: Boxplot representation of the distribution of entropy of model responses at each level of the hierarchy. The entropy tends to decrease as we move from level 1 to level 6.

Theorems & Definitions (4)

  • Proposition 3.1
  • proof
  • Proposition 3.2
  • proof