Overview:

A fixed-size model trained on a fixed-amount of data can gain predictive performance if the data is re-selected from a simpler distribution. Also, a model may learn dangerous capabilities from the data it is trained on. So, I recommend constraining language modeling pretraining data in order to build better and safer language models. Background, more details, findings, and next steps below.

Background:

I previously trained models on a constrained version of the TinyStories dataset (https://zenodo.org/records/23306194) and analyzed reasoning and generalization in this regime (https://zenodo.org/records/24075817). A ~10x smaller model trained on ~10x less data (constrained to a smaller world with less unique places, objects, adjectives, etc.) achieved comparable LLM-as-a-judge story output quality scores to the original TinyStories-33m param model, just within the smaller world.

More details:

I view pretraining data as being sampled from a distribution that could described by a theoretical, very complex, unrealized knowledge graph. I view entities (people, places, things, concepts, actions, adjectives, etc.) as nodes in this knowledge graph, and I view relationships (is-a, has-a, part-of, etc.) as being probabilistic connections between nodes. A knowledge graph can be simplified by pruning nodes; pretraining data can be constrained by deleting entities.

I create a whitelist of allowed lemmas (for example: "run") where Kuperman's age of acquisition (AoA) is (or probably would be) < 10. When filtering based on this whitelist, I allow all forms of a lemma (for example: "running", "runs", "ran", etc.). It is a simple filter... a more advanced filter would use named entity recognition (and a named entity whitelist) to identify multi-word entities, such as "large language model" (I am not using this type of filter for now due to my compute constraints).

I constrain both web-scale pretraining data (specifically the Nemotron-CC-HQ hq-paper split) and popular commonsense benchmarks (HellaSwag, WinoGrande, PIQA, and CommonsenseQA).

Findings:

I do two continual pretraining cooldown runs starting from an existing 1.7B model checkpoint (that has been trained for around ~240B tokens): one on data drawn from my filtered dataset and one on data drawn from the original dataset. I evaluate on both my filtered benchmarks and the original commonsense benchmarks.

The AoA<10 run matches or leads on performance on all 4 filtered benchmarks, and leads on 3 of the 4 unfiltered benchmarks.

(From my paper) Figure 3: Zero-shot accuracy over training tokens. Top: the items kept by my wordlist filter. Bottom: items removed by the wordlist filter. All runs start from the same checkpoint. Bands are ±1 standard error.

Paper: https://zenodo.org/records/24184396

Next steps:

Train an open 7B language model on 1T tokens (focused on commonsense reasoning, in-context learning, instruction following, and performing safe everyday tasks).