What should Baby Models read? Exploring Sample-Efficient Data Composition on Model Performance (2024)

Abstract: We explore the impact of pre-training data composition on the performance of small language models in a sample-efficient setting. Using datasets capped at 10 million words, we evaluate several data sources—including child-directed speech (CHILDES), classic fiction (Gutenberg), a mixed dataset (Mix), and synthetic TinyStories—across different model sizes ranging from 18 million to 705 million parameters. Our experiments show that smaller models (e.g., GPT2-18M and GPT2-44M) benefit from training on diverse datasets like Mix, achieving better performance on linguistic benchmarks. In contrast, larger models (e.g., GPT2-97M, GPT2-705M, and LLaMA-360M) perform better when trained on more complex and rich datasets like Gutenberg. Models trained on the CHILDES and TinyStories datasets underperformed across all model sizes. These findings suggest that the optimal dataset for sample-efficient training depends on the model size, and that neither child-directed speech nor simplified stories are optimal for small language models of all sizes. We highlight the importance of considering both dataset composition and model capacity for effective sample-efficient language model training.

## introduction in recent years, advancements in natural language processing have been largely driven by scaling language models to unprecedented sizes. various large-language model (llm) scaling laws have been formulated @xcite , with perhaps the most influential being the chinchilla law, which demonstrates that parameters and tokens scale approximately linearly as the model scales @xcite . many subsequent llms have been trained following this model @xcite , with some models including the llama 2 and llama 3 family of models being trained on 2 and 15 trillion tokens respectively, far more than the 'optimal' amount according to the chinchilla scaling law @xcite . however, it is often prohibitive to train such large models, and impractical to continue scaling with the amounts of data required to train such models. this has sparked interest in small language models @xcite with much fewer parameters, requiring much less data for training. while much research has been conducted on knowledge distillation and improving the model architecture for small language models, comparably less research has investigated the contributions of different types of data used for model training, which is arguably just as important. indeed, because llm pretraining data typically comprises a mix of sources @xcite , researchers have found that the composition of pretrained data greatly affects model performance @xcite , though determining the optimal recipe for pretraining data is challenging. recent research exploring optimization of pretraining data for llms at scale includes doremi, which trains a small proxy model to produce domain weights for downstream tasks, and then uses the model to resample the dataset for training huge llms @xcite . however, the question of how to choose data for sample-efficient training of small language models, such as in cases where computational resources are limited, has received little attention. psycholinguistic precedent exists for sampleefficient pretraining; children see much less words than a modern llm yet perform exceptionally well on reasoning tasks. for example, chinchilla sees over 10000 times the number of words a 13 year old child has ever encountered @xcite . by the time typical english-speaking children at around 6 years old have obtained adult-level grammatical knowledge @xcite , they have seen only around 10-50m words @xcite . in comparison, llama-3 is trained on 15t tokens @xcite . given the great disparity between the amount of training data an llm requires and what children require, it seems worthwhile to investigate whether training llms can be as sample efficient. babyberta @xcite attempts to address this, showing that when training a model on data similar to what is seen by children between the ages 1 and 6, it is able to acquire grammatical knowledge similar to pretrained robertabase, but with around 15x fewer parameters and 6,000x fewer words; this indicates that utilizing child-directed input may be advantageous for more sample efficient pretraining @xcite . similarly, @xcite follow suit, releasing tinystories, a synthetic dataset of short stories that only contain words that typical 3-to 4-yearold children understand. they demonstrate that tinystories can be leveraged to train language models with much less parameters than sota models, yet still produce coherent output with almost perfect grammar as well as emergent reasoning abilities. along the same vein, gpt-wee @xcite shows that child-directed speech can be used with curriculum learning for simulating children's learning as a potential solution to sample-constrained training. in this paper, we evaluate the effect of different datasets on model performance for sample efficient model training. in our case, we limit our training dataset to 10m words, in accordance with the babylm challenge's super-strict track @xcite . we consider several different types of datasets, namely child-directed speech (childes), classic books (gutenberg), a mixed dataset (mix) and the tinystories dataset. experimental results show that smaller models benefit from training on diverse datasets like mix on the babylm evaluation suite @xcite , but larger models perform better when trained on more complex and rich datasets like gutenberg. our findings suggest that the optimal dataset depends on the model size and that neither childdirected speech nor child-directed stories are optimal for language models of any sizes. ## dataset for our experiments, we obtained datasets from the babylm challenge @xcite . individual categories of 10m-word datasets were procured by extracting the first 10m words from that cate-gory in the 100m-word dataset of the babylm challenge. we also used mix, the 10m-word developmentally-plausible corpus of babylm, and tinystories. to measure for complexity in the language of these datasets, we use several readability metrics, including the flesch reading ease (fre) score @xcite , ari (automated readability index) @xcite , and the gunning fog index @xcite . for a document d @xmath0, its fre score is computed as: higher ari scores indicate more complex text requiring higher grade levels to comprehend. the gunning fog index score is calculated as: our individual datasets comprise: ## results and discussion overall, our results demonstrate that the effectiveness of a training dataset is dependent on the model size. specifically, smaller models (with fewer parameters) benefit more from training on a diverse on the other hand, we see a very different story when looking at macro average glue scores for the models ( ## limitations our study has several limitations. first, we used consistent hyper-parameters across all experiments for comparability, but this may not have been optimal for each model-dataset pair. tuning hyperparameters individually could have yielded better performance. second, the blimp and ewok benchmark assess linguistic competence on tasks on represented in datasets such as tinystories or childes, potentially biasing the evaluation. in short, there is a mismatch between the training data afforded by child datasets and the test set. lastly, due to computational limitations, models were trained for only four epochs. longer training might have allowed models to better capture the nuances of the datasets. ## conclusion and future work in this paper, we investigated the impact of dataset composition on the performance of small language models in a sample-efficient training regime. by training models of varying sizes on different datasets limited to 10 million words, we sought to identify which types of data are most beneficial for language acquisition in resource-constrained settings. we found that tiny models (e.g., gpt2-18m and gpt2-44m) performed best when trained on the mix dataset, which offers a diverse combination of language inputs, while slightly larger small language models achieved superior performance when trained on the gutenberg dataset, leveraging its linguistic richness. in contrast, models trained on childes or tinystories underperformed regardless of size. for future work, a more thorough investigation of other types of data sources such as news articles, scientific texts, and conversational data might better tease out the optimal dataset for model performance. additionally, it might be useful to explore curriculum learning, which presumable models the developmental process of a language learning child. widening the benchmarks beyond glue and blimp tasks to coherent text generation, as well as scaling dataset sizes and tasks would allow for a more comprehensive and robust study as well.