NeurIPS 2025

Analyzing Similarity Metrics for Data Selection for
Language Model Pretraining
Dylan Sam∗
Carnegie Mellon University
Ayan Chakrabarti
Google Research
Afshin Rostamizadeh
Google Research
Srikumar Ramalingam
Google Research
Gui Citovsky
Google Research
Sanjiv Kumar
Google Research
Abstract
Measuring similarity between training examples is critical for curating high-quality
and diverse pretraining datasets for language models. However, similarity is
typically computed with a generic off-the-shelf embedding model that has been
trained for tasks such as retrieval. Whether these embedding-based similarity
metrics are well-suited for pretraining data selection remains largely unexplored.
In this paper, we propose a new framework to assess the suitability of a similarity
metric specifically for data curation in language model pretraining applications.
Our framework’s first evaluation criterion captures how well distances reflect
generalization in pretraining loss between different training examples. Next, we use
each embedding model to guide a standard diversity-based data curation algorithm
and measure its utility by pretraining a language model on the selected data and
evaluating downstream task performance. Finally, we evaluate the capabilities of
embeddings to distinguish between examples from different data sources. With
these evaluations, we demonstrate that standard off-the-shelf embedding models
are not well-suited for the pretraining data curation setting, underperforming even
remarkably simple embeddings that are extracted from models trained on the same
pretraining corpus. Our experiments are performed on the Pile, for pretraining
a 1.7B parameter language model on 200B tokens. We believe our analysis and
evaluation framework serves as a foundation for the future design of embeddings
that specifically reason about similarity in pretraining datasets.
1
Introduction
The recent success of language models [Brown et al., 2020, Chowdhery et al., 2023] is in no small
pa...