Research map / Large language models / Language, safety and interpretability
Analyzing Similarity Metrics for Data Selection for Language Model Pretraining
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes an evaluation framework (loss-variance reduction within clusters, diversity-based curation with a 1.7B model on 200B Pile tokens, and source cluster purity) showing off-the-shelf embeddings underperform simple embeddings from small LMs trained on the same corpus.
Problem
Off-the-shelf text embedding models are used to measure similarity for pretraining data curation, but their suitability for this purpose is unexplored and hard to evaluate at pretraining scale.
Approach
A three-part evaluation framework for embeddings in pretraining data curation: (1) variance reduction of pretraining loss within embedding-based clusters, (2) downstream performance of a 1.7B LM trained on diversity-curated subsets, and (3) cluster purity w.r.t. Pile data sources. Specialized embeddings extracted from a small LM trained on the same corpus (LM Token Embeds, LM Output Embeds) are sh
Key results
- LM Output Embeds achieves the highest average downstream score (53.8 ± 0.1) vs. 51.9 for random baseline and 53.3 for Gecko.
- LM Token Embeds (no forward pass) matches Gecko (53.1 vs 53.3 average) and outperforms USE, BERT among off-the-shelf models on average.
- Specialized embeddings from the small trained LM yield higher variance reduction than off-the-shelf models across all cluster sizes and pretraining checkpoints.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use diversity-based curation instead of random subset selection | Pile pretraining data selection · Pile downstream tasks (Average) | average task performance: 51.9 $\pm$ 0.1 -> 53.8 $\pm$ 0.1 | Helped |
| Use LM Output Embeds instead of Gecko embeddings for diversity curation | Pile pretraining data selection · Pile downstream tasks (Average) | average task performance: 53.3 $\pm$ 0.3 -> 53.8 $\pm$ 0.1 | Helped |
| Use specialized LM embeddings rather than off-the-shelf embeddings for loss-based clustering | Pile pretraining loss clustering · Pile | variance reduction | Helped |
| Use LM Token Embeds for diversity curation instead of off-the-shelf embeddings | Pile pretraining data selection · Pile downstream tasks | average task performance | Helped |
| Increase the number of PCA components before clustering (ablation) | Gecko and LM Output Embeds · Pile pretraining loss clustering | variance reduction | Helped |
| Replace PCA with random projections for dimensionality reduction (ablation) | Gecko and LM Output Embeds · Pile pretraining loss clustering | variance reduction | Hurt |
| Replace PCA with random projections for dimensionality reduction (ablation) | BERT, USE, and LM Token Embeds · Pile pretraining loss clustering | variance reduction | Hurt |
Surprising findings
- Gecko, a newer retrieval-focused embedding, does not improve over older USE and BERT on pretraining loss variance reduction.
- A no-forward-pass averaged token embedding matches or outperforms off-the-shelf embedding models in downstream curation.
Limitations
- Empirical demonstration restricted to the Pile dataset.
- Simplified diversity-based curation algorithm used to isolate embedding quality; results may not transfer to more sophisticated curation methods.
- Reported diversity-curation differences between many methods are small (often within ~1 point).
Experimental setup
Datasets: The Pile
Benchmarks: ARC BoolQ SuperGLUE HellaSwag MultiRC OpenBookQA PiQA RACE ReCoRD RTE Story Cloze WiC
Metrics: variance reduction cluster purity top-1 accuracy F1
Baselines: Universal Sentence Encoder Gecko BERT Random subset selection
Builds on: RAC (reciprocal agglomerative clustering) UL2
Closest papers by meaning
- SimReg: Achieving Higher Performance in the Pretraining via Embedding Similarity Regularization (NeurIPS 2026 · similarity 0.90)
- Repurposing Language Models into Embedding Models: Finding the Compute-Optimal Recipe (NeurIPS 2024 · similarity 0.89)
- Predictive Data Selection: The Data That Predicts Is the Data That Teaches (ICML 2025 · similarity 0.89)
- DEPT: Decoupled Embeddings for Pre-training Language Models (ICLR 2025 · similarity 0.89)
- MATES: Model-Aware Data Selection for Efficient Pretraining with Data Influence Models (NeurIPS 2024 · similarity 0.89)
- Towards Scalable Data Diversification for Language Model Pretraining via Leverage Score Sampling (NeurIPS 2026 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.