Analyzing Similarity Metrics for Data Selection for Language Model Pretraining

TL;DR

This paper evaluates the effectiveness of existing similarity metrics for selecting training data in language model pretraining, revealing that standard off-the-shelf embeddings are often inadequate for this purpose.

Contribution

It introduces a new framework to assess and compare the suitability of similarity metrics specifically for pretraining data curation in language models.

Findings

Standard off-the-shelf embeddings underperform for data curation in pretraining.

Simple embeddings from models trained on the same corpus can be more effective.

The framework helps guide the development of better similarity metrics for pretraining datasets.

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.