Target-Aware Language Modeling via Granular Data Sampling (2024)

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 selecting 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 ~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.

## introduction language model pretraining is the cornerstone of universal language models (lms), creating generalpurpose representations to excel across a variety of nlp downstream tasks @xcite . this process often involves the use of vast amounts of text, sometimes measured in billions or even trillions of tokens from webpages @xcite . however, there are instances where a model needs to perform well in specific domains while not compromising performance in others. this necessitates the use of data selection methods to determine which potential data points should be included in the training dataset and how to effectively sample from these selected points @xcite . subword: multi-granular tokenization token is ization crucial tokenization is crucial tokenization is crucial word: multi-word: tokenization token is crucial ization n-gram table multi-granular: importance weights 1 doc2vec 2 data sampling 3 token ization is crucial tokenization tokenization is 0 1 2 3 4 5 29 23 13 8 17 9 hashed features is crucial 32 21 adding exclude tokenization is crucial for enabling efficient handling of textual information ... ... document k highest importance scores one approach to reducing data size is coreset selection, which involves selecting a small, representative subset of data @xcite . coresets can significantly decrease computational costs while maintaining robust performance. in this work, we explore the optimization of coresets towards a target data distribution, but relaxing the data sampling process to reduce domain biases. here we revisit importance sampling @xcite by proposing to utilize tokens of different granularities as features, ranging from subword, word, to multi-word (or n-gram) tokens (shown in figure 1 ). we observed empirically that by controlling the granularity of tokens in the tokenizer, we can construct coresets with less domain biases -fine-grained tokens capture more task knowledge while coarse-grained tokens preserve general information. thus, we experiment with adapting the vocabulary set of pretrained tokenizers to data from specific tasks, and modulating token granularity in the vocabulary to maintain generality. to demonstrate the efficacy of multi-granular sampling, we use eight downstream tasks as target tasks with llama-3's tokenizer being the base tokenizer, where its vocabulary set is adapted to task data. this process creates a domain-specific vocabulary set which we use to featurize text documents 1 for target-aware data sampling with higher quality sampled data. specifically, we demonstrate that these pretrained models, using only about 1% of the data, perform on par with the full refinedweb data and outperform randomly selected samples across eight benchmarks ( §4). this contributes to the ongoing discourse on optimizing pretraining data to improve computational efficiency and model performance. unlike past approaches which can be computationally-intensive @xcite @xcite or simple but easily biased towards target data distribution @xcite , our approach is unique in two significant ways: ## the approach selecting samples from large-scale datasets such as refinedweb @xcite is slow and expensive. a tractable solution is to encode each document as a vector using n-gram features that can be computed easily. here we assume a small number of target text examples d task from a target distribution p and a large raw dataset d raw from a distribution q with n examples. we aim to select k examples (k ≪ n ) from the raw dataset that are similar to the target. we adopted the importance sampling technique as in xie et al. ( 2023 align with target distribution. the technique provides a tractable importance estimate of each text and applies importance sampling on a feature space z that provides the necessary structure. the feature extractor h ∶ @xmath0. the resulting raw and target feature distributions are q feat and p feat , respectively. our objective is to select examples whose features align with the target feature distribution p feat . to do so, features q feat and p feat are extracted (figure 1 ) using n-grams extracted from each tokenized document using an adapted tokenizer. each n-gram is mapped to a key in the hash table where the ids of the table define a fixedsize embedding, and each key maps to the n-gram count. then, the importance weights are computed for each featurized example z @xmath1, with the weight w @xmath2. the final step involves sampling, where we select k examples without replacement from a categorical distribution, the probabilities of which are given by ## experimental setup network, training details and evaluation. we pretrain the decoder-only transformer using causal language modeling objectives on selected datasets, averaging over three initialization runs for each configuration, where model weights were randomly initialized. the language models varied in size, with 125m, 350m, 500m, and 1.5b parameters. this range allowed us to explore how model complexity impacts the final results. pretraining was conducted on a distributed computing setup with 32 gpus across 4 nodes, each equipped with an h100 graphics card. coreset selection. we evaluated our proposed multi-granular selection approach against random selection (random) and compared it with the same sampling algorithm using word-based n-gram features. importance sampling (xie et al., 2023) was employed for all feature types. 125m 350m 500m 1.5b model size 45.6 47.6 49.6 51.6 53.6 55.6 57.6 accuracy (%) multi-granular n-gram random datasets. we evaluate the models on eight common sense reasoning tasks in a zero-shot fashion, including arc-easy, arc-challenge @xcite , boolq @xcite , piqa (bisk et al., 2020), siqa @xcite , hel-laswag @xcite , obqa @xcite , and winogrande @xcite . ## main results overall, language models varied in four sizes display marked improvements when trained on sampled coresets selected using multi-granular features, achieving a 6.94% improvement in average benchmark scores (avg.), as shown in ## further discussion finer-grained features reduce task biases. based on our ablation, we observe marked improvement by simply using subword n-grams. moreover, we show in figure 4 that selecting from a single task introduces task data biases that degrades the performance. this is mitigated through the use of finer-grained n-gram features where we introduce multi-granular tokens containing subwords and multi-words, which gives an additional 5.78% gram and random baseline across eight tasks using 125m models, trained solely on data selected with arc-easy data as the target. relative performance is used. we observe that multi-granular features enable the model to consistently outperform the baseline despite task-specific biases in the data. we performed similar experiments for all other benchmarks in the appendix, where for all eight tasks, the same pattern is observed that multi-granular n-grams yield almost no degradation across benchmarks. ## conclusion and future works in this study, we revisited importance sampling of text corpus in language modeling by exploring multi-granular n-grams as features. this led us to explore a pretraining paradigm where we can obtain more targeted data for more efficient lan-guage modeling. our findings, validated across eight benchmarks, allow us to put forward multigranular n-grams features as viable document representations used in importance sampling. for future work, we will aim to extend this approach to larger language models and datasets.