Efficient Training of Language Models with Compact and Consistent Next Token Distributions (2024)

Abstract: Maximizing the likelihood of the next token is an established, statistically sound objective for pre-training language models. In this paper we show that we can train better models faster by pre-aggregating the corpus with a collapsed n -gram distribution. Previous studies have proposed corpus-level n -gram statistics as a regularizer; however, the construction and querying of such n -grams, if done naively, prove to be costly and significantly impede training speed, thereby limiting their application in modern large language model pre-training.We introduce an alternative compact representation of the next token distribution that, in expectation, aligns with the complete n -gram distribution while markedly reducing variance across mini-batches compared to the standard next-token loss. Empirically, we demonstrate that both the n -gram regularized model and our approximation yield substantial improvements in model quality and convergence rate compared to existing methods. Furthermore, our approximation facilitates scalability of gains to larger datasets and models compared to the straightforward n -gram regularization method.

## introduction since the advent of the first neural language models (nlm) @xcite , a standard approach to training nlms has been to maximize the likelihood of the next token (nt) given the preceding tokens in randomly sampled token sequence from the dataset. in this paper we show that llm pre-training can be made significantly more efficient by supervising with a distribution over multiple possible next tokens, instead of a single next token. earlier studies @xcite @xcite have explored using corpus-level n-gram statistics to improve the quality of rnn based language models. this can also be thought of as label-smoothing i ) by querying an n-gram model which can be slow during training. our proposed method, coconts, truncates and approximates the y all i and stores the preprocessed distribution (y cc i ) along with the dataset itself for faster retrieval during training. @xcite with n-gram estimated distribution. in this paper we show that regularizing with corpus level n-gram statistics continues to be beneficial even on current language models, both when pre-training from scratch or fine-tuning. however, a major hurdle is scaling up such techniques to today's corpus and model size. we address these challenges by (1) proposing a compact representation of the next token distribution while continuing to be statistically consistent. and (2) designing data structures to efficiently make use of these distributions without stalling throughput-optimized linear algebra accelerators (xla devices) such as tpus. in particular, we find that our proposed data handling strategy and objective (named "coconts", pronounced "coconuts") often matches the performance of n-gram models with full (non-truncated) distribution (sec. 5.2). notably, while the prior n-gram augmented methods use storage proportional to the size of training dataset, our storage overheads are constant with respect to size of the training dataset (sec. 5.3). we also find that ngram augmented training methods (prior works as well as coconts) can reach the same validation perplexity as the ntl objective in nearly 50% less optimization steps (sec. 5.3). our proposal is useful for improving model quality in both pre-training (sec. 5.6) as well as fine-tuning -all parameter (sec. 5.2) or parameter efficient (sec. 5.5). our key contributions can be summarized as: (1) we highlight the usefulness of training with n-gram statistics for faster convergence of language models, and propose a statistically consistent truncation strategy of making their implementation scalable and practical for current data and model sizes (sec. 4). (2) we show how this truncated distribution can be efficiently retrieved during training time without slowing down xla devices (sec. 4.1, sec. 4.2). ## related work language modeling. early neural language models proposed in @xcite and @xcite were trained to maximize the likelihood of the next token in randomly sampled training batches. while the transformer architecture @xcite has been the powerhouse of modern llms @xcite @xcite , the key training objective has remained the same. several studies have explored alternate objectives such as unlikelihood training @xcite or contrastive loss @xcite to improve text generation from these lms. however, these objectives are not statistically consistent, and have not been adopted in large scale pre-training. augmenting language model training. some prior work @xcite @xcite have proposed extending rnn based lms @xcite with kn smoothed n-gram lms, and shown that these result in a better language model. this sparked interest @xcite @xcite to introduce corpus level n-gram statistics into an otherwise local training procedure. @xcite replaced the standard cross-entropy loss with a ranking based loss for which they used n-grams as a weak teacher to obtain "gold" rankings. @xcite explore additional kl divergence penalty similar to @xcite and our work. however, they obtain ground truth distribution by similarity between word embeddings trained on the corpus. our work is closest to the n-gram regularized loss of @xcite but we propose a modified compact supervision that scales to large corpus. we also compare our models against @xcite to study whether kl divergence from learned word vectors can serve as a better proxy for relatively expensive count based n-gram models. ## language modeling let v denote a vocabulary of tokens. given a token sequence t : t 1 , t 2 . . . , t i where each t j in an integer index into the vocabulary v, a language model p θ (t|t 1 . . . t n ) ∈ r |v| assigns a multinomial distribution of the probability of the possible next token that could follow t. the size of the token sequence t is limited to a maximum length l. for training θ we are given a corpus d @xmath0, . . . , x n where n is typically very large n ≫ l. a standard method of training is to maximize the likelihood of the next token (ntl or nt) given the preceding tokens in a randomly sampled snippet of length l from d train . let x j , . . . , x j+l denote such a snippet sampled at position j of d train . the training objective then becomes in this paper we investigate if the above convergence to the corpus-level next-token distribution can be sped up via changing the training objective to directly match the target distribution y all t . for long prefixes t we do not expect too much repetition, and also maintaining the y all t proportions for all possible prefixes may incur too much overhead. so, we fix a maximum prefix length k, and instead optimize for a mixture of these two objectives. the allnts objective. the above loss is reminiscent of the use of the corpus-level n-gram statistics to regularize lm training @xcite . the second term goes over all n-grams t of length from 1 to k, and attempts to match the learned model distribution to the observed fraction of next tokens in the corpus following that n-gram t. ## compact and consistent next tokens: coconts we propose to approximate the full empirical next token distribution y all x j:n with a more compact and consistent supervision y cc let y topr x j:n denote a truncation of the y all x j:n where only the top r largest fractions are retained and the rest of the truncated to zero. let 1 v (t i ) be one-hot encoding of size |v| with component for t i as 1. we approximate y all choosing u, v. we choose u, v such that e(y cc this shows that we are "stealing" some probability mass from the top-r token positions and assigning them to the rare token positions, so that across all repetitions of the prefix we have consistent supervision on the next token distribution. by choosing u carefully we can control this consistency. the value of u has to be in the range [0, 1/(1p)] for y cc to be non-zero for all positions. the distance |y cc x j:ny all x j:n | can be computed as as u increases from zero, the supervision at the rare token positions is made closer to the ideal, but at the cost of the frequent token positions. we define a hyper-parameter γ > 1, and choose @xmath1. while p varies with prefixes x j:n and datasets, we found that a value of γ = 1.5 performs well across diverse settings. we show effectiveness of the approximation via an example in appendix a.2. ## experiments through the experiments, we empirically validate the effectiveness of our approximation (sec. 5.2) along with the ablation studies on various hyperparameters. we also evaluate the training efficiency of coconts in sec. 5.3. finally, we present two case studies on coconts in pre-training (sec. 5.6) and parameter-efficient fine-tuning (sec. 5.5) to further demonstrate the usefulness and relevance of coconts even in modern llm usecases. hyperparameters for all the experiments are presented in appendix a.1. ## conclusion in this work, we revisited the benefits of regularizing language model training with corpus-level n-gram statistics, and proposed ways to scale up their implementation on current scales of data and model sizes. our proposal truncates the n-gram estimated next-token distribution and introduces a novel method of mixing with occurrences of frequent and rare tokens so as to provide low-variance supervision of the distribution of next tokens. the distributions are designed to be compact and can be stored with the corpus so that their retrieval is as simple as a disk read operation. we empirically show that coconts performs comparable to allnts objective but is significantly more efficient than allnts. notably, while the allnts storage costs scale with dataset, coconts storage costs depend only on the hyperparameters k and r. our fine-tuning experiments suggest that coconts based training benefits smaller models the most with larger models seeing only modest improvements. we also observe that imposing all-nts or coconts loss only on a small k token prefix was sufficient to improve the overall model performance. case study on the babylm challenge highlights that coconts trained base lms are better than standard ntl trained base lms on downstream tasks as well.