Analyzing Similarity Metrics for Data Selection for Language Model Pretraining

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Sam, Dylan, Chakrabarti, Ayan, Rostamizadeh, Afshin, Ramalingam, Srikumar, Citovsky, Gui, Kumar, Sanjiv
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866911222253748224
author Sam, Dylan
Chakrabarti, Ayan
Rostamizadeh, Afshin
Ramalingam, Srikumar
Citovsky, Gui
Kumar, Sanjiv
author_facet Sam, Dylan
Chakrabarti, Ayan
Rostamizadeh, Afshin
Ramalingam, Srikumar
Citovsky, Gui
Kumar, Sanjiv
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2502_02494
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Analyzing Similarity Metrics for Data Selection for Language Model Pretraining
Sam, Dylan
Chakrabarti, Ayan
Rostamizadeh, Afshin
Ramalingam, Srikumar
Citovsky, Gui
Kumar, Sanjiv
Machine Learning
Computation and Language
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.
title Analyzing Similarity Metrics for Data Selection for Language Model Pretraining
topic Machine Learning
Computation and Language
url https://arxiv.org/abs/2502.02494