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.

text: Findings of the Association for Computational Linguistics ACL 2024, pages 12051–12064 August 11-16, 2024 ©2024 Association for Computational Linguistics
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.

heading: 1 Introduction
text: Since the advent of the first neural language models (NLM) (Bengio et al., 2003; Mikolov et al., 2013), 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 (Neubig and Dyer, 2016; Zhao et al., 2017; Yang et al., 2019) 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

heading: Corpus level
text: (Szegedy et al., 2016; Müller et al., 2019) 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”,
12051
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). (3) We discuss practical strategies and their effects on model performance of scaling this strategy to potentially very large scale datasets via sharding (Sec. 5.3). (4) We thoroughly test our proposal through ablation studies, comparisons with prior works on both fine-tuning as well as pre-training of language models (Sec. 5).

heading: 2 Related Work
text: Language modeling. Early neural language models proposed in (Bengio et al., 2003) and (Mikolov et al., 2013) were trained to maximize the likelihood of the next token in randomly sampled training batches. While the Transformer architecture (Vaswani et al., 2017) has been the powerhouse of modern LLMs (Black et al., 2022; Zhang et al., 2022; Touvron et al., 2023), the key training objective has remained the same. Several studies have explored alternate objectives such as unlikelihood training (Welleck et al., 2020) or contrastive loss (Su et al., 2022; Jain et al., 2023) 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 (Mikolov et al., 2011; Chelba et al., 2014; Jozefowicz et al., 2016) have proposed extending RNN based LMs (Bengio et al., 2003;
Mikolov et al., 2013) with KN smoothed n-gram LMs, and shown that these result in a better language model. This sparked interest (Neubig and Dyer, 2016; Zhao et al., 2017; Yang et al., 2019) to introduce corpus level n-gram statistics into an otherwise local training procedure. Frydenlund et al. (2022) 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. Li et al. (2020) explore additional KL divergence penalty similar to Yang et al. (2019) 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 Yang et al. (2019) but we propose a modified compact supervision that scales to large corpus. We also compare our models against (Li et al., 2020) to study whether KL divergence from learned word vectors can serve as a better proxy for relatively expensive count based n-gram models.

heading: 3 Language Modeling
text: Let V denote a vocabulary of tokens. Given a token sequence t : t1, t2 . . . , ti where each tj in an integer index into the vocabulary V , a language model Pθ(t|t1 . . . tn) ∈ 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 Dtrain = x1, . . . , xN 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 Dtrain. Let xj , . . . , xj+L denote such a snippet sampled at position j of Dtrain. The training objective then becomes
θ̂D = argmaxθ
L∑
n=1
logPθ(xj+n+1|xj , . . . , xj+n)
(1) When a prefix xj:n = xj , . . . , xj+n has multiple occurrences in the corpus, then for the same context depending on the sampling position j, the target may be different. But the model Pθ needs to converge to a consistent distribution. For example, a prefix “The United” could occur multiple times within a corpus with different possible next tokens as shown in Figure 1. If t = t1, . . . , tn denotes the tokens in such a prefix, it can be seen that at conver-
gence, the next token likelihood Pθ(t|t1, . . . , tn) for a token sequence is expected to be equal to the empirical distribution of tokens following t over the entire corpus Dtrain. Let yallt denote the fractional frequency of the occurrence of tokens of V following all different positions where t appears in the corpus Dtrain. At convergence we expect:
Pθ̂D(t|t) −→ y all t , where,
yallt [w] =
∑ j∈D δ(xj:n = t, xj+n+1 = w)∑
j∈D δ(xj:n = t)
(2)
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 yallt . For long prefixes t we do not expect too much repetition, and also maintaining the yallt 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.

heading: The AllNTs objective.
text: min θ
L∑
n=k+1
− logPθ(xj+n+1|xj:n)
+
k∑
n=1
KL (yallxj:n ;Pθ(·|xj:n)) (3)
The above loss is reminiscent of the use of the corpus-level n-gram statistics to regularize LM training (Yang et al., 2019; Neubig and Dyer, 2016). 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.
Benefits of AllNTs. When we supervise the model to match empirical distribution on all possible next tokens after a prefix, the convergence of the model is expected to be faster. In the empirical section we will show that training with even small n-grams (k = 4), gives rise to much lower perplexity for the same computation budget than the original training only for next token likelihood (Eq 1).
Overheads of AllNTs. For imposing the AllNTs loss we need to create a data structure like a trie, which for each possible prefix can provide the distribution of next tokens (Heafield, 2011; Heafield et al., 2013). Querying the trie for every sampled mini-batch is inefficient. These inefficiencies are especially noticeable (Sec. 5.3) when scaling to
larger datasets (>1B tokens) as the batch creation (which includes trie lookup) on CPU is slow. In the next section we show an alternative method for supervising the next token distribution that significantly reduces these overheads.

heading: 4 Compact and Consistent Next Tokens: CoCoNTs
text: We propose to approximate the full empirical next token distribution yallxj:n with a more compact and consistent supervision yCCxj:n at each sampled prefix xj:n. Unlike yallxj:n which can be of size as large as the vocabulary size, the alternative we propose is of size at most r + 1 where r is a chosen hyperparameter, like 4 or 8 in our experiments. We design yCCxj:n so that in expectation over the minibatches, yCCxj:n matches the AllNTs supervision but where the variance across loss terms is significantly smaller than via the supervision in the NT objective.
Let ytoprxj:n denote a truncation of the y all xj:n where only the top r largest fractions are retained and the rest of the truncated to zero. Let 1V(ti) be one-hot encoding of size |V| with component for ti as 1. We approximate yallxj:n with y CC xj:n as follows:
yCCxj:n =
{ vy
topr xj:n if xj+n+1 ∈ ytoprxj:n uy topr xj:n + 1V(xj+n+1) otherwise
(4)
Choosing u, v. We choose u, v such that E(yCCxj:n) ≈ yallxj:n and |yCCxj:n −yallxj:n | is minimized. Let p = ∑ t y
topr[t] denote the total probability mass covered by the top-r highest probability tokens. An example appears in Figure 1. For a token t ∈ V that is outside the top-r tokens in yallxj:n , it is clear that E(yCCxj:n)[t] = y all xj:n [t]. Now consider a token t in the top-r set. We want to determine values of u, v s.t. E(yCCxj:n)[t] = y all xj:n [t]
E(yCCxj:n [t]) = y topr[t](vp+ u(1− p)) = ytopr[t]
=⇒ v = 1− (1− p)u p
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/(1− p)] for
yCC to be non-zero for all positions. The distance |yCCxj:n − yallxj:n | can be computed as
|yCCxj:n−yallxj:n | { = (1− p)u if xj+n+1 ∈ ytoprxj:n ≤ 2− up otherwise
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 u = 1 γ−p . While p varies with prefixes xj: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.
4.1 Pre-enriching the dataset with yCC
The above design of yCC allows us to tackle the core computational inefficiency of AllNTs. Instead of incurring the CPU overheads of trie-lookups during training, we propose to pre-enrich the training corpus with the compact yCC distributions stored along with the corpus. The total storage cost becomes only (L+2kr)/L times the original storage cost but we avoid the expensive trie lookup operations during training. Also, we can build the entire trie in-memory only once during the pre-processing phase and discard after the data enrichment phase.
We implement our trie in C, similar to Heafield (2011); Heafield et al. (2013). Each TrieNode consists of the count and a HashMap where key is the next token ID and value is a pointer to the next TrieNode. The HashMap is implemented using an AVL tree. We implement a Top-r query method which returns Top-r token IDs and Top-r probabilities for a given prefix of upto k tokens using Eq. 2.Concretely, we first construct the trie by reading sequences of k tokens from the dataset and inserting it in the trie. At every level i, we increment the count by 1 to implicitly record prefix [t1 . . . ti]. Once the trie is constructed, we start reading (disjoint) sequences of L tokens from the dataset and writing the “enriched” sequence back to disk. To “enrich” a sequence xj , . . . , xj+L, we first look up the prefix in the trie and get pointers to k nodes i.e. one at each level. At each of these nodes, the HashMap is storing yall. We can efficiently traverse this HashMap in descending order of fractions to get ytopr. Once we get Top-r(yallj:i) ∀i ∈ [0, k), we can write the “enriched” sequence to disk. Note that, this operation still requires disk storage of (L+ (L+ 2kr))×N tokens where N is the total
number of sequences. We present additional discussion and a sample walkthrough of this procedure for more clarity in Appendix A.3.
4.2 Building the mini-batch with yCCi To build the mini-batch of size B in standard NTL, one simply reads sequences of L tokens B times and concatenates them. In standard NTL, the sequence of L tokens itself is both input and target but in CoCoNTs, we have 2 targets that need to be built. In CoCoNTs we need to read sequence of L + 2kr tokens from the disk and use the first L tokens to form x,y similar to NTL. Next 2r “tokens” correspond to ytopr where the first token is the actual token ID and the next token is really the count of that token appearing after prefix [t1]. The 2r tokens following this would correspond to Top-r(yall2 ) which encodes counts and tokens appearing after prefix [t1, t2]. This would repeat for k times to give k distributions as targets for KL divergence in Eq. 3. We provide additional discussion on memory footprint in Appendix A.4.

heading: 5 Experiments
text: 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.

heading: 5.1 Datasets, Baselines and Metrics
text: We explore the effectiveness of CoCoNTs in finetuning existing models on WikiText-103 (Merity et al., 2017), MiniPile (Kaddour, 2023; Gao et al., 2020) and a subset of PubMed-Abstracts (Luo et al., 2022). The WikiText-103 training split consists of ≈ 114M tokens while MiniPile and PubMed splits consist of ≈ 1.6B and 2.6B tokens respectively.
For our full fine-tuning experiments, we compare CoCoNTs against AllNTs (Yang et al., 2019) and NTL objectives on WikiText-103. We compare against D2GPO baseline (Li et al., 2020) as well since they also use a KL divergence based augmentation of training loss.
Since Su and Collier (2023) raised concerns
about isotropy of the base gpt2-125m (Radford et al., 2019) model, we also study effects of CoCoNTs objective on gpt-neo-125m (Black et al., 2021) and opt-125m, opt-1.3B (Zhang et al., 2022).
Metrics. Following (Su et al., 2022), we evaluate each fine-tuning method several on model quality and text quality metrics. • Perplexity (ppl) of the model on the test set. • Prediction accuracy (acc) of the model. Given
a sample with inputs [x1 . . .xL] and labels [y1 . . .yL] from the test set, we take argmax of each of the L predicted distribution at time step t to get top-1 predicted token and match it against yt to calculate the prediction accuracy. • Repetition (rep) measures the fraction of top-1 next-token predictions that occur in the prefix. • Expected calibration error (ECE) measures how over/underconfident is the model when making correct/incorrect predictions. Lower ECE indicates better calibrated models. • MAUVE (MVE) (Pillutla et al., 2021) measures the similarity between the generated text and reference text using the embeddings of another large pretrained model. • Repetition within a generated single text sequence: (rep-n) 100× (1− |unique n-grams||total n-grams| ). • Diversity (div.) measures repetition at different n-gram levels: ∏4 n=2(1− rep-n100 ). • Number of unique bigrams (#uniq) in the generated text. We also compare the Zipf coefficient of the gen-
erated text to gold text as suggested by (Meister et al., 2023). For generating text, we either use greedy decoding or nucleus (Holtzman et al., 2020) sampling. To compare training efficiency, we use following metrics: • Number of optimization steps (steps-to-ppl)
taken to reach NTL’s perplexity on the val set. • Total wall clock time (TWT) to finish the entire
training. We exclude the time for preprocessing (trie building and storage) as the preprocessing is a one-time operation which many times did not take very long. Additionally, one could always start with datasets that are preprocessed by someone else. We do include the time it takes to load the trie and training time retrieval as these operations will often stall the XLA devices. • Total disk usage (disk) required for training (includes storage of prefix trie). • Maximum CPU RAM (max-RAM) used for pretraining i.e. loading and using the trie.

heading: 5.2 Model performance
text: CoCoNTs is comparable to AllNTs and better than NTL. In Table 1, we compare performance of CoCoNTs with various other objectives. We find that CoCoNTs is able to provide consistent gains over the NTL baseline. We also notice that CoCoNTs outperforms D2GPO (Li et al., 2020) which indicates that count based conditional n-gram models are able to provide a stronger training signal as compared to word embedding similarity. Maximum gains with CoCoNTs are observed on gpt2125m, however this could be related to isotropy of gpt2-125m checkpoint as discussed by (Su and Collier, 2023). Across all metrics and models, AllNTs is generally the best performing model while CoCoNTs comes close to it. CoCoNTs offers gains (albeit modest as compared to small models) to larger models (opt-1.3B) as well.

heading: 5.3 Training efficiency
text: CoCoNTs is significantly more efficient than AllNTs. Figure 2 compares our training efficiency metrics across various training methods on WikiText-103 and PubMed datasets. AllNTs uses Python’s defaultdict 1 to implement the trie HashMap while CoCoNTs uses our efficient implementation of HashMap as described in Sec. 4.1. AllNTs serializes the resultant trie to disk using the Python’s pickle2 library. We also explore using an existing n-gram implementation (Heafield, 2011) with Python FFI as denoted by AllNTs-CFFI. Because AllNTs stores the n-gram model on disk for later retrieval, its disk overhead grows with the size of the dataset as opposed CoCoNTs’s constant (dependant on k, r only) disk overhead. Measuring the total wall clock time, we find that CoCoNTs (k = 4) is faster than AllNTs (k = 2) which highlights the benefits of our approximation and preenriching of the dataset.
Results on PubMed demonstrate challenges in scaling up vanilla AllNTs to large datasets. On our machine with 256GB RAM, the prefix trie for k = 8 with vanilla AllNTs did not fit in memory. It is possible that AllNTs-CFFI is able to fit everything in memory similar to CoCoNTs but we could not explore in depth since our n-gram implementation (Heafield, 2011) kept crashing on our system. While AllNTs-CFFI saves disk and RAM,
1 https://docs.python.org/3/library/collections.html#collections.
defaultdict 2 https://docs.python.org/3/library/pickle.html
it still is not as efficient as CoCoNTs on WikiText103. Moreover, the trie loading time as well as the (somewhat) slow CFFI interface serialization overheads significantly increase total wall time of AllNTs-CFFI over CoCoNTs. Effects of sharding. As compared to datasets used in this work, modern LLMs (Black et al., 2022; Touvron et al., 2023; Groeneveld et al., 2024) are often trained on far bigger web corpora (Gao et al., 2020; Together Computer, 2023; Soldaini et al., 2024) for which building an n-gram model in-memory may not be feasible. In such cases, we show that such datasets can be sharded into multiple small datasets of several few billion tokens with each shard being enriched with its own n-gram index. We study the effect of sharding by simulating it on MiniPile and PubMed datasets and seeing effect on perplexity as shown in Fig. 3. We find that after a certain threshold of shards, the number of tokens per shard decreases so much that KL penalty can become overly optimistic and result in worse perplexity. In general, we found that having more than billion tokens per shard was sufficient to get results close to AllNTs while still using modest amount of RAM.

heading: 5.4 Understanding CoCoNTs
text: Ablation studies on k and r. By default we choose k = 8, r = 8. We report ablations on these values when fine-tuning gpt2-125m on WikiText-103 and
compare using validation perplexity. As indicated by trends in Fig. 4, we find that increasing k and r leads to predictable improvements in perplexity as compared to the NTL baseline. We do notice a significant difference in perplexities between AllNTs and CoCoNTs for lower values of k. This could be potentially due to poorer yCCi approximation since the fan-out is expected to be significantly higher for initial few tokens. Empirically, in the first few levels of the trie, the branching factor is the highest often leading to yalli that has a support size much larger than r. This is further supported by the significant improvement observed when going from r = 2 to r = 4 for a fixed k. Moreover, as k increases, the support of yalli naturally decreases. In fact for k > 2, the average support size is less than 4 on WikiText-103. This implies that p = 1 in Eq. 4 leading to v = 1 which effectively reduces CoCoNTs to AllNTs objective.
Using larger/domain specific data sources to build n-gram models can help. We show that using n-gram statistics from alternative corpus is also useful. To study this, we augmented our existing WikiText-103 n-gram trie with n-grams from MiniPile and PubMed datasets independently. As shown in Table 2, model trained with WikiText103 + MiniPile index improves perplexity on both WikiText-103 and MiniPile. On the other hand, if the augmenting dataset (PubMed) is both large and
Wiki-103 MiniPile PubMed
gpt2-125m 24.279 43.456 37.819 AllNTs 20.634 41.261 38.538 CoCoNTs 20.717 42.486 39.016
+MiniPile 20.223 36.899 38.564 +PubMed 20.926 39.432 36.175
Table 2: Perplexity improvements when augmenting n-gram index with larger/domain specific data. Each model (row) is trained on WikiText-103 and evaluated on validation splits of (column) WikiText-103, MiniPile and PubMed. CoCoNTs with indices built on additional MiniPile or PubMed data improves performance on respective datasets.
domain specific, the resultant model improves on augmenting dataset at the cost of performance on the original (WikiText-103) dataset.
Perplexity reduction as a function of prefix length. We show that even though we impose the
CoCoNTs loss only in a small prefix of an overall sequence of length 256, the improvement in model quality is throughout the length of the sequence. For this, we measure the negative log-likelihood (NLL) separately at each position from 1 to 256 in the test data and show our findings in Figure 5.We observe that we reduce NLL at all sequence positions over the two existing methods. This suggests that in today’s NLM architectures with shared parameters, enforcing the correct loss in only a subset of the positions is sufficient for overall gains.
Effect of longer pretraining. We study that pretraining for more steps flattens the benefits of CoCoNTs slightly but not completely. We continued pretraining of gpt-neo-125m from our experiments (Sec. 5) for 20k more steps (50% more than original pretraining steps) and report our findings in Table 3. Looking at the validation perplexities, we do see that the difference between baseline and CoCoNTs perplexities does decrease from 2.5 to 2.3
20 21 22 23
k
21
22
23 24 P er p le xi ty NTL
20 21 22 23 24
NTL
CoCoNTsAllNTs CoCoNTs
21 23 25 27
Position in sequence
3
4
5
6
N eg
at iv
e lo
g lik
el ih
oo d
NTL SimCTG CoCoNTs (k = 2) AllNTs (k = 2) CoCoNTs (k = 8) AllNTs (k = 8)
over the 20k additional steps. To correctly assess if this difference goes to zero, significantly longer training is required.

heading: 5.5 Case Study: PEFT Domain Adaptation
text: Parameter efficient fine-tuning (PEFT) of LLMs (Hu et al., 2022; Dettmers et al., 2023) has become a popular method to adapt general purpose LLMs such as LLaMA-2 (Touvron et al., 2023) on a specific domain. To study benefits of CoCoNTs objective in this direction, we first train a LoRA (Hu et al., 2022) to fine-tune LLaMA-2-7B on our
split of the PubMed dataset using each of the finetuning method. Both AllNTs and CoCoNTs use k = 4 while CoCoNTs uses r = 8. With this LoRA as the starting point, we fine-tune the resultant (Medical LLaMA) model on 2 medical QA (PubMedQA (Jin et al., 2019) and MedQA (Jin et al., 2021)) tasks independently. As summarized by accuracy of these QA tasks in Table 4, we find results similar to full fine-tuning i.e. CoCoNTs performs better than NTL but worse than AllNTs. However, AllNTs incurs an almost 60% increase in pre-training time, whereas our methods reduces the overheads by half as seen by the last total wallclock time (TWT) column in the table.

heading: 5.6 Case Study: The BabyLM Challenge
text: Can CoCoNTs be used to pre-train a better language model for downstream tasks? We pre-train an opt-125m architecture model from scratch on the data from the BabyLM challenge (Warstadt et al., 2023). We use the nearly 100M word data from the “strict” track with standard preprocessing. Once all (NTL, AllNTs, CoCoNTs) the models are pre-trained, we finetune each on 3 downstream tasks (with recommended hyperparameters) from the evaluation suite – CoLA (Warstadt et al., 2019), MRPC (Dolan and Brockett, 2005) and RTE (Dzikovska et al., 2013) which are subsets of the SuperGLUE benchmark (Wang et al., 2019b,a) for the BabyLM challenge. Task specific fine-tuning does not use any custom objective. Table 5 show metrics on respective tasks. Both AllNTs and CoCoNTs trained base LMs show significant improvement over standard NTL in CoLA which is about linguistic acceptability. This could be due to grammatically incorrect sentences/completions automatically being suppressed in the n-gram index. On other downstream tasks such as paraphrase detection (MRPC) or entailment (RTE), the performance of all models is much closer in comparison. Also, observe that the total wall clock time (TWT) for AllNTs is 35% higher than baseline, and CoCoNTs reduces the overhead down to 19%.

heading: 6 Conclusion
text: 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.

heading: Limitations and Future Work
text: In this section, we discuss the limitations of CoCoNTs objective and provide insights into potential challenges and areas for improvement. The onetime preprocessing step required by our method for very large-scale datasets requires sharding. Effects of such sharding on social biases of the model must be studied carefully. When applying our method to very large datasets like C4 or The Pile, the implementation of the prefix-trie using better optimized libraries as discussed in (Jurafsky and Martin, 2023) may become necessary. This could require significant engineering efforts to optimize access times and ensure efficient training. To address scalability concerns further, a possible suggestion in addition to sharding would be to “sparsify” the trie for such large-scale datasets. By pruning branches with low counts, we can significantly reduce the overall memory footprint while still maintaining the essence of (idea of “heavy hitters” (Misra and Gries, 1982; Woodruff, 2016; Braverman et al., 2017)) the next token distribution.
When applying our method to very large datasets like C4 or The Pile, the implementation of the prefix-trie using better optimized libraries as discussed in (Jurafsky and Martin, 2023) may become necessary. This could require significant engineering efforts to optimize access times and ensure efficient training. To address scalability concerns further, a possible suggestion in addition to sharding would be to “sparsify” the trie for such large-scale datasets. By pruning branches with low counts, we can significantly reduce the overall memory footprint while still maintaining the essence of the empirical next token.
Furthermore, it is important to acknowledge that our CoCoNTs objective aims to match the empirical next-token distribution, and thus inherits any biases present in the training data. However, an advantage of our approach is that the prefix trie allows for detailed exploration and identification of these biases. If such biases are observed, it should be possible to edit the prefix trie to mitigate their influence. Our experimental setup included the
largest models and datasets that could run comfortably on our compute resources. Future work can explore effectiveness of CoCoNTs on even larger datasets (e.g. ThePile (Gao et al., 2020), RedPajama (Together Computer, 2023), Dolma (Soldaini et al., 2024) etc.) and larger scale models such as LLaMA2 (Touvron et al., 2023).

heading: Ethics Statement
text: We acknowledge that our objective entails preprocessing and handling large-scale datasets for creating the prefix trie. This necessitates careful attention to privacy concerns and the implementation of robust data protection measures. It is vital to thoroughly examine and mitigate any biases that may emerge in the training data prior to the application of our proposed objective.
Moreover, given the improved text generation capabilities demonstrated by our approach, it is imperative to address ethical considerations regarding the responsible use of language models trained using our proposed objective. In this context, we underscore the significance of ensuring that the deployment and utilization of such models align with ethical standards, including but not limited to mitigating the potential for malicious use, promoting fairness in algorithmic decision-making, and safeguarding user privacy.

heading: A Appendix

heading: A.1 Hyperparameters
text: Most of our experiments are performed on a single TPUv2-8 core VM with TPU VM architecture. We also sometimes used 4x NVIDIA A100 GPUs with FlashAttention for a fraction of all experiments.
For fine-tuning experiments (Sec. 5.2), we start with publicly available checkpoints for gpt2125m3, gpt-neo-125m4, opt-125m5 and opt-1.3B6. Each 125m parameter model is trained for 40k steps with AdamW (Loshchilov and Hutter, 2019) optimizer with effective batch size of 192. The maximum learning rate was set to 10−4 with 10% of maximum steps as warmup followed by cosine decay to zero. For the 1.3B parameter model, we set the maximum steps to 10k and reduce the batch size to 32. Each fine-tuning run took roughly 5-6 hours on TPUs and 8-10 hours on GPUs.
For pre-training on the BabyLM challenge (Sec. 5.6), we set the batch size, optimizer and learning rate schedule similar to fine-tuning and trained for total of 5 epochs. For downstream tasks, we use the hyperparameters mentioned in the BabyLM evaluation pipeline7. The pre-training took 5.5 hrs on TPUs while fine-tuning on downstream tasks took 1-1.5hrs each.
For parameter-efficient fine-tuning of LLaMA (Sec. 5.5), we set the batch size to 32 and finetune on PubMed with hyperparameters similar to WikiText-103 fine-tuning. We use the same parameters as BabyLM downstream tasks for LLaMA PEFT downstream tasks as well. The pre-training took close to 6 hours on TPUs while fine-tuning on downstream tasks took 2-2.5hrs each.
A.2 Effectiveness of yCC
Example. We show a simple example to illustrate how yCC manages to reduce variance in the next-token distribution across sampled minibatches. Assume vocab size is 5, and yallt = [0.6, 0.3, 0.1, 0.05, 0.05]. Let r = 2. Assume γ = 1.5. For this case we have p = 0.9, u = 1.66, v = 0.5. Everytime we sample a mini-batch where next token is from the top-2 set: {1, 2}, we supervise with yCC = v[0.6, 0.3, 0, 0, 0] which has a distance of 0.166 from yall. Whereas, when we
3 https://huggingface.co/gpt2 4 https://huggingface.co/EleutherAI/gpt-neo-125m 5 https://huggingface.co/facebook/opt-125m 6 https://huggingface.co/facebook/opt-1.3B 7 https://github.com/babylm/evaluation-pipeline
sample the last three tokens we are at most 0.51 from yall. Contrast this to the baseline NT case where for rare token the distance to yall could be as high as 1-0.05=0.95! With yCC we reduce this distance to 0.51. Even for frequent tokens the distance has been reduced from a maximum of 0.7 to 0.166.
Effect of r. A crucial hyperparameter in our approximation is r. While we see the effects of r on overall model quality in Sec. 5.4, we study the effect of r in a more controlled way when learning a single 10-class multinomial. By increasing r, as shown in Fig. 6, we find that both convergence rate as well as KL divergence between learned and actual multinomial consistently improves.

heading: A.3 Additional Discussion on Pre-enriching the Dataset
text: In this section, we provide a step-by-step walkthrough of Pre-enriching process for the dataset.
First, we assume that the prefix-trie is already created and available in-memory for k = 2. This trie has a crucial property: the “count” attribute (64 bit unsigned integer in our implementation) of each node Ni indicates how many times the prefix, which corresponds to the path from the root to Ni, appears in the training dataset Dtrain. The root node stores the count of total prefixes in the dataset. Figure 7 shows an example trie, where the node associated with the token “United” has a count of 1000, while the root node has a count of 40000. This means that there are 1000 sequences in total
that begin with the word “United”. Additionally, the child node “States” has a count of 526, indicating that there are 526 sequences that start with “United States”.
Let’s assume that a block (length L) beginning with "United States of America . . . ” is selected in the pre-enriching second pass over the dataset. Since k = 2, we retrieve from the prefix trie with prefix “United States” as shown in Figure 7. At every level (i.e. highlighted Ni), first yalli will be created in floating point representation with the same bitwidth as token IDs. In our case, this was fp16. Then, we sort the distribution to get top r token IDs and top r probability values. Storing these on disk will require space equivalent to 2kr more tokens.
After a block of L tokens is read from the input file, it is immediately written as is to the output file. Then we find top r token IDs and probabilities for each k and sequentially write these values to the output file. After all k such distributions are written, we would have written equivalent of L+2kr tokens to the output file.
A.4 Additional Discussion on Minibatch building with yCCi
Since we know that the maximum support of all the yCCi distributions is r+1, we can ideally easily pass them as key-value pairs to the training loop and calculate KL divergence more efficiently. This causes only O(kr) increase in memory footprint of a batch. However, this can be slightly inefficient since we need to run “gather” operation to obtain correct components of the predicted distribution
Pθ(yi) based on token indices. “gather” operation is often slow8 on TPU/XLA devices which rely on a predictable dataflow in order to optimize their compute graph. Prior works such as BigBird (Zaheer et al., 2020) have resorted to special resizing and converted the operation to continuous memory access.
Such tricks are harder to implement here without ascertaining an upper bound on the maximum token ID in the support of yCCi . Ideally, obtaining such bounds may be useful and possible since the tokenizer (such as WordPiece or BPE) are expected to assign lower token IDs (earlier “merges”) to frequent tokens anyway. In our implementation, we initialize a k × |V| size zero vector and use the “scatter” 9 operation to populate counts at correct places. While the “scatter” operation is also slow, we perform it during batch creation on CPU which is latency optimized as opposed to previous proposal which was doing “gather” operation on accelerators which are throughput optimized. While this increases the memory footprint of the batch by O(k|V|), we found that using such dense vectors for KL divergence resulted in the model running slightly faster on both TPUs and GPUs.
Since we pack the [yCC1 , . . . ,y CC k ] directly into the batch as |V|-dimensional vectors, the memory used by the labels in the CoCoNTs objective is considerably higher than in the baseline. Despite this, we occupy only 0.5% of the total TPU/GPU RAM used by the trainer. The baseline (NT) method takes only 128 × 256 × 2 = 64KB of memory to store labels for a sequence length of 256 and a batch size of 128, assuming16-bit integer token IDs. In the CoCoNTs objective, we provide additional k distributions. Assuming each number in the distribution is a 16-bit float, and considering k = 8 and GPT2’s vocabulary size of 50257, we occupy approximately 128×8×50257×2 = 98.16MB of additional GPU/TPU RAM per batch during training. On our GPUs, the NT objective utilizes around 72GB of RAM out of the total available 80GB, leaving more than enough room to accommodate all the k = 8 extra distributions per sequence per batch.
8 https://github.com/pytorch/xla/issues/898, https://github.com/
pytorch/xla/issues/3587 9 https://pytorch.org/docs/stable/generated/torch.scatter.html

heading: Future Work
text: Future work can explore effectiveness of CoCoNTs on even larger datasets (e. g. ThePile (Gao et al. , 2020), RedPajama (Together Computer, 2023), Dolma (Soldaini et al. , 2024) etc. ) and larger scale models such as LLaMA2 (Touvron et al. , 2023)