Target-Aware Language Modeling via Granular Data Sampling

TL;DR

This paper introduces a data sampling method using multi-granular n-gram features to efficiently pretrain language models focused on specific domains, achieving comparable or better performance with significantly less data.

Contribution

It proposes a novel importance sampling approach with multi-granular n-gram features for targeted domain-specific language model pretraining.

Findings

Models trained on ~1% of data match full data performance.

Sampled data correlates strongly with downstream task success.

Outperforms random sampling across multiple benchmarks.

Abstract

Language model pretraining generally targets a broad range of use cases and incorporates data from diverse sources. However, there are instances where we desire a model that excels in specific areas without markedly compromising performance in other areas. A cost-effective and straightforward approach is sampling with low-dimensional data features, which allows to select large-scale pretraining data for domain-specific use cases. In this work, we revisit importance sampling with n-gram features consisting of multi-granular tokens, which strikes a good balance between sentence compression and representation capabilities. We observed the sampled data to have a high correlation with the target downstream task performance while preserving its effectiveness on other tasks. This leads to the proposed data sampling paradigm where language models can be pretrained more efficiently on selected documents. On eight benchmarks we demonstrate with $\sim$1% of the data, pretrained models perform on par with the full RefinedWeb data and outperform randomly selected samples for model sizes ranging from 125M to 1.5B.