Analysing The Impact of Sequence Composition on Language Model Pre-Training (2024)

Abstract: Most language model pre-training frameworks concatenate multiple documents into fixedlength sequences and use causal masking to compute the likelihood of each token given its context; this strategy is widely adopted due to its simplicity and efficiency. However, to this day, the influence of the pre-training sequence composition strategy on the generalisation properties of the model remains underexplored. In this work, we find that applying causal masking can lead to the inclusion of distracting information from previous documents during pre-training, which negatively impacts the performance of the models on language modelling and downstream tasks. In intra-document causal masking, the likelihood of each token is only conditioned on the previous tokens in the same document, eliminating potential distracting information from previous documents and significantly improving performance. Furthermore, we find that concatenating related documents can reduce some potential distractions during pre-training, and our proposed efficient retrieval-based sequence construction method, BM25Chunk, can improve incontext learning (+11.6%), knowledge memorisation (+9.8%), and context utilisation (+7.2%) abilities of language models without sacrificing efficiency.

heading: 1 Introduction
text: Large Language Models (LLMs) are pre-trained on large amounts of documents by optimising a language modelling objective and show an intriguing ability to solve a variety of downstream NLP tasks (Brown et al., 2020; Biderman et al., 2023; Touvron et al., 2023; Jiang et al., 2023). Previous works emphasise the importance of pre-training data quality (e.g., Gunasekar et al., 2023; Lee et al., 2022; Tirumala et al., 2023; Soboleva et al., 2023) and diversity (e.g., Xie et al., 2023; Gao et al., 2021;
https://github.com/yuzhaouoe/pretraining-data-packing
Kaddour, 2023) to improve the generalisation properties of language models. However, the influence of the pre-training sequence composition strategy remains largely under-explored.
For most decoder-only language model pretraining pipelines (e.g., Shoeybi et al., 2019; Ott et al., 2019; Brown et al., 2020; Biderman et al., 2023; Geng, 2023; Liu et al., 2023b; Zhang et al., 2024), constructing a pre-training instance involves packing, which refers to the process of combining randomly sampled documents into a chunk that matches the size of the context window; and causal masking, which refers to predicting the next token conditioned on all previous tokens, including those from different documents in the chunk. An alternative to causal masking is intra-document causal masking, where the likelihood of each token is conditioned on the previous tokens from the same document; intra-document causal masking is not commonly used in existing open-source pretraining frameworks as it is argued to adversely impact pre-training efficiency (Brown et al., 2020; Pagliardini et al., 2023). However, to the best of our knowledge, there is no systematic analysis in the literature on how causal masking affects the generalisation properties of models despite its role in improving efficiency.
To analyse the impact of the packing and masking strategies on pre-training, we pre-train language models using intra-document causal masking (referred to as INTRADoc, Section 2.2) and compare them with models pre-trained via causal masking with several packing strategies by varying the relatedness of the documents in the pre-training chunks. Specifically, we analyse the results produced by a commonly used baseline method that randomly samples and packs documents (MIXChunk); a method that samples and packs documents from the same source based on their meta-information
7897
(UNIChunk); and our proposed efficient retrievalbased packing method, which retrieves and packs related documents (BM25Chunk, Section 2.1).
Our experimental results indicate that using causal masking without considering the boundaries of documents can lead to the inclusion of distracting information from previous documents during pre-training (Section 3 and Section 5.1), negatively impacting the performance of the models in downstream tasks (Section 4). We observe that intradocument causal masking, which eliminates the potential distractions from irrelevant documents during pre-training, can significantly improve the performance of the model while increasing its runtime (+4% in our implementation, see Appendix A).
We also find that improving the relatedness of the documents in pre-training chunks can reduce some potential distractions from previous documents (e.g., UNIChunk avoids packing documents from different distributions, such as code and news text), which can improve the performance of causal masking models on a wide array of downstream tasks. Finally, we show that our proposed efficient retrieval-based packing method, BM25Chunk, can improve a model’s language modelling (+6.8%), in-context learning (+11.6%), knowledge memorisation (+9.8%), and context utilisation (+7.2%) abilities using causal masking and thus without sacrificing pre-training efficiency.
Our main contributions and findings can be summarised as follows: • We systematically analyse and compare the mod-
els pre-trained using causal masking and intradocument causal masking; our experimental results reveal that using causal masking without considering the boundaries of documents can result in significant performance degradation (Section 3 and Section 4). • We find that improving the relatedness of the documents in each pre-training chunk benefits causal masking models, and our proposed efficient retrieval-based packing method (BM25Chunk, Section 2.1) can improve the performance of language models significantly. • We quantitatively analyse the attention distribution of the models during language modelling (Section 5.1), and investigate the burstiness property of pre-training chunks (Section 5.2); our findings indicate that models can be more robust to irrelevant contexts and obtain better performance when improving the relatedness of documents in pre-training chunks.

heading: 2 Packing and Masking Strategies for Pre-Training Sequence Composition
text: In this section, we formally introduce the pretraining data packing strategies, causal masking, and intra-document causal masking.

heading: 2.1 Packing Strategies
text: Let Di represent a corpus, such as Wikipedia, C4, or GitHub, and let D = ⋃sDs denote the dataset resulting from the union of such corpora. Furthermore, each corpus Ds is defined as a set of documents Ds = {d1, . . . , d|Ds|} , where each document di is defined as a sequence of tokens di = ( x1, . . . , x|di| ) .
A packing strategy involves first selecting a set of documents {di}ni=1 from D, and then packing them into a chunk C with a fixed length |C| = L. Following Brown et al. (2020), we concatenate the documents {di}ni=1 by interleaving them with endof-sentence ([EOS]) tokens to construct a chunk. A packed sequence (or chunk) C is denoted as:
C = (d1[EOS]d2[EOS] . . . SPLIT(dn)), (1)
where [EOS] is the end-of-sentence token, SPLIT() truncates the last document such that |C| = L, and the content of the chunk C will be removed from the dataset D to avoid sampling the same documents multiple times.
In the following, we introduce three strategies to sample the documents {di}ni=1 from the dataset D for composing each pre-training chunk, namely MIXChunk, UNIChunk, and BM25Chunk.
MIXChunk In MIXChunk (baseline), documents di ∈ D are sampled uniformly at random from the entire pre-training corpus D:
di ∼ Uniform(D).
As a result, in MIXChunk, a chunk can contain documents from different source datasets, e.g., Wikipedia and GitHub, as shown in Figure 1(a).
UNIChunk In UNIChunk, each chunk is composed of documents from a single source corpus Ds:
di ∼ Uniform(Ds), with Ds ⊆ D.
This helps to avoid packing documents from different distributions (such as code and news text) together. To construct a training batch, we sample sequences from each corpus Ds in proportion to the number of tokens in Ds.
2
Wikipedia
Github
ArXiv
doc-4doc-3
doc-3doc-1 doc-2
UniChunk
MixChunk
(a) MIXChunk randomly samples documents from all corpora to construct pre-training sequences, which can pack documents from different sources. UNIChunk randomly samples documents from a single source to construct a sequence.
Document Buffer
doc-2
doc-1
Step-1: start from doc-1
doc-3
doc-1 doc-2
Step-2: retrieve from the buffer using doc-1
Step-3: return doc-2
Step-4: retrieve from the buffer using doc-2
Step-5: return doc-3
BM25Chunk
Step-6: return the remaining part to the buffer
(b) The sequence construction process in BM25Chunk. The left part represents a document buffer that caches a set of documents randomly sampled from the corpus.
Figure 1: Packing strategies for pre-training chunks construction. (a) illustrates the compositions of MIXChunk and UNIChunk; (b) presents the sequence construction process of BM25Chunk.
BM25Chunk To improve the relevance of documents in pre-training chunks, we employ a BM25based retriever to construct pre-training chunks, referred to as BM25Chunk. Specifically, given a document di ∈ Ds, we retrieve a sequence of documents {di}ni=1 by di+1 = RETRIEVE(di,Ds); here, RETRIEVE(di,Ds) retrieves the most similar documents to di from Ds based on BM25 scoring.
However, this retrieval process can be computationally inefficient due to the size of the pretraining corpus Ds. To improve the efficiency of the retrieval step, we restrict the retrieval scope to a subset Bs ⊆ Ds of the corpus Ds, reducing the computational complexity of retrieval; the proposed approach is outlined in Figure 1(b). More formally, we introduce a document buffer Bs ⊆ Ds that contains k documents uniformly sampled from
Ds, which serves as the retrieval source for constructing pre-training chunks:
d1 ∼ Uniform(Bs), di+1 = RETRIEVE(di,Bs).
After retrieving a sequence of documents {di}ni=1 from the buffer Bs for constructing a chunk, we refill the buffer by sampling new documents from Ds. The time complexity analysis and more details are presented in Appendix C.

heading: 2.2 Masking Strategies
text: Another core element of LLM pre-training is the masking strategy, which determines how nexttoken prediction distributions are conditioned on other tokens in the sequence.
Causal Masking In causal masking, each token in a sequence is predicted solely based on all preceding tokens in the sequence. More formally, given a chunk C = (x1, . . . , x|C|) defined as in Equation (1), the likelihood of C is given by:
P (C) =
|C|∏
i=1
P (xi | x1, . . . , xi−1),
where P (xi | x1, . . . , xi−1) denotes the probability of the token xi given all preceding tokens x1, . . . , xi−1 in the chunk. During pre-training, causal masking implies that, given a chunk C, the probability of each token in C will be conditioned on all preceding tokens, including those belonging to different documents. Causal masking is the standard practice when pre-training decoder-only LLMs (e.g., Shoeybi et al., 2019; Brown et al., 2020; Zhang et al., 2022; Biderman et al., 2023; Geng, 2023; Liu et al., 2023b; Zhang et al., 2024).
Intra-Document Causal Masking In intradocument causal masking, on the other hand, the probability of each token is conditioned on the previous tokens within the same document. More formally, given a chunk C defined as in Equation (1), the probability of each token dij belonging to document di is only conditioned on the preceding tokens within di:
P (C) =
n∏
i=1
|di|∏
j
P ( dij | di1, . . . , di(j−1) ) .
We refer to models trained using intra-document causal masking as INTRADoc. The details of our efficient implementation using FlashAttention (Dao, 2023) are available in Appendix A.
3

heading: 3 Language Model Pre-Training

heading: 3.1 Settings
text: Pre-Training Corpora In this work, we use SlimPajama (Soboleva et al., 2023) as the pretraining corpus, which consists of seven subcorpora, including CommonCrawl, C4, Wikipedia, GitHub, StackExchange, ArXiv, and Book. This allows us to investigate packing strategies in a mixed corpora setting. We sample documents with 150B tokens from SlimPajama as the pre-training corpus and ensure each subset maintains the same proportion of tokens as in the original dataset.
Pre-Training Models The model implementation is based on the LLaMA (Touvron et al., 2023) architecture with minor modifications to support intra-document causal masking. We pre-train 1.3B parameters models using context windows of 2,048 (referred to as 2K) and 8,192 (8K) tokens. We use the same set of documents with the difference in pre-training sequence composition to pretrain models, including causal masking models, i.e., MIXChunk, UNIChunk, and BM25Chunk, and intra-document causal masking models INTRADoc. More details are available in Appendix B.
Previous works (Brown et al., 2020; Pagliardini et al., 2023) argued that dynamic sequence-specific sparse masking reduces training efficiency. Compared to causal masking, we observe a 4.0% efficiency degradation on intra-document causal masking in our implementation, and the discussion on implementation is presented in Appendix A.

heading: 3.2 Results
text: For evaluating LLMs trained under different packing strategies, in this work, we compute the perplexity (PPL) of a held-out set of documents where each document is treated independently. The results are summarised in Table 1.
We can see that BM25Chunk achieves the lowest PPL among the three causal masking models, yielding a lower average PPL compared to MIXChunk in the 2K (−0.62) and 8K (−0.78) settings. Furthermore, UNIChunk also yields a lower average PPL than the baseline MIXChunk (−0.34 and −0.21). These results indicate that increasing the relatedness of documents in a sequence can improve the language modelling ability of models.
When considering models trained via intradocument causal masking, we can see that INTRADoc achieves the lowest PPL compared to all models trained via causal masking. This indicates eliminating the potential distracting information from irrelevant documents during pretraining benefits the language modelling ability of models. Specifically, we observe that both BM25Chunk and INTRADoc obtain significantly lower PPLs on GitHub, where INTRADoc improves over UNIChunk in both the 2K (−1.3 PPL) and 8K (−2.0) models. For UNIChunk, though we avoided packing web text and code, its improvement over MIXChunk on GitHub is slight. This phenomenon could imply that code pre-training is more adversely affected by the distraction of unrelated context, and both intra-document causal masking and retrieval-based sequence construction strategy can alleviate this issue.

heading: 4 Experiments on Downstream Tasks
text: In the following, we evaluate the in-context learning, knowledge memorisation, and context utilisation abilities of the models.

heading: 4.1 In-Context Learning
text: Following Shi et al. (2023), we evaluate in-context learning abilities of the models using seven text classification datasets, namely SST2 (Socher et al., 2013), Amazon (Zhang et al., 2015), Yelp (Zhang
4
et al., 2015), DBpedia (Lehmann et al., 2015), AGNews (Zhang et al., 2015), and TweetEval hate/offensive tweet detection tasks (Barbieri et al., 2020).
In Table 2, we report the in-context learning accuracy values of the models in few-shots learning settings, using 20 and 48 demonstrations for 2K and 8K models, respectively. We truncate the input sequences to fit within their respective context windows. For models pre-trained using causal masking, we can see that UNIChunk produces more accurate results than MIXChunk, while BM25Chunk yields a higher average accuracy than MIXChunk for 2K (+11.6%) and 8K (+11.3%) models. These results indicate that increasing relatedness of the documents in pre-training chunks can improve the in-context learning abilities of the models.
In Figure 2, we present the average accuracy using different numbers of few-shot demonstrations. We observe that BM25Chunk has an on-par accuracy with INTRADoc on the 2K setting; however, INTRADoc obtains a significantly higher accuracy compared to BM25Chunk on the 8K setting. It may imply that using a longer context window size can result in increased distractions for causal masking pre-training; meanwhile, constrained by the performance of the retrieval method, BM25Chunk
decreases the accuracy on the 8K setting. For 8K models, MIXChunk and UNIChunk obtain similar results to their corresponding 2K models, and they do not improve the accuracy when increasing the number of demonstrations. It might be due to the similar levels of distraction in both 2K and 8K settings using random packing strategies.

heading: 4.2 Knowledge Memorisation
text: We use two open-domain question-answering (ODQA) datasets, namely NaturalQuestions (NQ, Kwiatkowski et al., 2019) and TriviaQA (TQA, Joshi et al., 2017), to evaluate the knowledge memorisation properties of the models. We use 12 demonstrations for the 2K models and 48 demonstrations for the 8K models. In Table 3, we show the mean Exact Match (EM) scores calculated based on 5 different sets of demonstrations.
For models trained with causal masking, we can see that increasing the relatedness of documents in pre-training chunks can improve the knowledge memorisation ability of models. Compared to the baseline MIXChunk, BM25Chunk obtains +9.8% and +3.0% EM improvements on 2K and 8K models, respectively. We also note that intra-document causal masking significantly improves the knowledge memorisation ability, especially for 8K models, where INTRADoc improves EM by +12.3%
5
and +37.5% over MIXChunk for 2K and 8K models, respectively. These results support our hypothesis that reducing the distractions deriving from concatenating multiple, potentially unrelated documents in pre-training chunks can improve the knowledge memorisation ability of the models.

heading: 4.3 Reading Comprehension and Retrieval-Augmented Generation
text: We evaluate the pre-trained models on a set of reading comprehension tasks, namely RACE (Lai et al., 2017), SQuAD (Rajpurkar et al., 2016), HotpotQA (Yang et al., 2018), and the following retrieval-augmented generation (RAG) tasks: NQ, TQA, and Multi-Document Question-Answering (MDQA, Liu et al., 2023a). For NQ and TQA, we use the top two passages retrieved by the dense retriever (Karpukhin et al., 2020; Izacard and Grave, 2021), denoted as NQ-open and TQA-open. Our results for RACE, SQuAD, and RAG tasks are summarised in Table 4, while the results on MDQA are available in Figure 3.
We can see that BM25Chunk produces more accurate results than MIXChunk and UNIChunk in all tasks and obtains the best average accuracy, showing that increasing the relatedness of documents in pre-training chunks can improve the context utilisation ability. Specifically, BM25Chunk obtains a significantly better accuracy on multi-hop QA task HotpotQA, showing it can better utilise multiple relevant information from the context.
INTRADoc obtains the best average accuracy in the 2K models and obtains the best scores in 5 of 6 tasks in the 8K models. It indicates that eliminating potential distractions from unrelated documents and learning each document independently can improve context utilisation ability. This finding is different from the ideas in previous works, which suggested that pre-training with multiple documents in one context (Shi et al., 2023) and adding distraction
in context during pre-training (Tworkowski et al., 2023) benefit context utilisation ability.
In MDQA, for each question, there are 30 documents provided in the context, where only one of them contains the answer to the question — MDQA is used to evaluate the ability of models to filter out irrelevant information and identify the relevant parts of a long context. This task has been used to analyse the lost-in-the-middle phenomenon in LLMs where they struggle to retrieve information stored in the middle of long contexts (Liu et al., 2023a). In the following, we analyse how the accuracy of models varies with the position of relevant information in the context. In these experiments, we focus on 8K models due to their ability to handle long contexts. The zero-shot results on MDQA are outlined in Figure 3. We observe that both BM25Chunk and INTRADoc tend to produce more accurate predictions than MIXChunk and UNIChunk when the relevant passage is located at the beginning or middle of the context. These results show that BM25Chunk and INTRADoc can better filter irrelevant context and locate relevant information; these results are further corroborated
6
by our experiments in Section 5.1 where we analyse the attention distribution of the models during the language modelling process.

heading: 5 Discussion and Analysis

heading: 5.1 Can Models Ignore Irrelevant Contexts Before the End-of-Sequence Token?
text: In the following, we analyse whether models can filter irrelevant context during language modelling by examining the attention score distributions over the context. Specifically, we concatenate two randomly sampled documents from the SlimPajama validation set, separate them by an end-of-sequence token [EOS], and check to which extent the attention distributions of the model focus on the irrelevant document in the sequence. More formally, we define the distraction proportion of the token in position p in the current document at layer l as:
DISTRPROP(l, p) = |Cd|∑
i=1
alp,i (2)
where |Cd| denotes the number of tokens in the irrelevant document, alp,i is the average multi-head attention scores to the i-th token in the irrelevant document Cd at layer l, and ∑|Cd|+p i=1 a l p,i = 1. In our experiments, we set |Cd| = 256, and the results are outlined in Figure 4.
We can see that the latter positions have lower distraction proportions but remain 45%-52% average distraction proportion until the 256th token of the second document, as shown in Figure 4(c). We find that models trained via BM25Chunk (green line) tend to have lower distraction proportions than other causal masking models, showing that they can better recognise relevant information in the context,
matching the results in Figure 3. The above analysis also demonstrates that during the pre-training, causal masking models can be distracted by unrelated documents in context, and the models can be more robust to irrelevant contexts when reducing distractions in pre-training sequences.
Furthermore, in Figure 4(d), we compare INTRADoc and causal masking models using "\n" as the separator instead of [EOS], because [EOS] can only appear at the end of sequences during pre-training using intra-document causal masking. The results indicate that INTRADoc has the lowest distraction proportion compared to causal masking models; meanwhile, BM25Chunk consistently has a lower distraction proportion than MIXChunk and UNIChunk using "\n" as the separator. These results match the finding in Section 4.3, where INTRADoc and BM25Chunk can better recognise relevant information in the context.

heading: 5.2 Burstiness Property of Sequences
text: Chan et al. (2022); Han et al. (2023) found a positive correlation between the model’s in-context learning ability and burstiness property of the training sequences. Here, burstiness refers to the phenomenon where certain types of tokens occur in clusters or bursts rather than being uniformly distributed across all documents. Burstiness is an inherent property of text; for example, a specific medical term might be frequently used in medical articles and rarely appear in general texts. Higher burstiness results in a lower Zipf’s coefficient of token frequency within a sequence (Han et al., 2023).
Following Han et al. (2023), we use Zipf’s coefficient to measure the burstiness property of pretraining sequences. Formally, let r denote the rank
7
of a token in a sequence, and f is a frequency function that maps the rank r to the frequency of that token in the sequence. Then, according to Zipf’s law, we have that f(r;α) ∝ 1rα , where α ∈ R+ is the Zipf’s coefficient; a lower α presents an increased burstiness property within the sequence.
In Table 5, we show the Zipf’s coefficients α on different pre-training sequences. Our results show that, for causal masking approaches that use the same chunk size, a lower Zipf’s coefficient, which denotes increased burstiness property, often results in more accurate results. However, INTRADoc can obtain significantly better results than UNIChunk with the same Zipf’s coefficient. The above results indicate that, for causal masking approaches, the correlation between higher burstiness and better performance could derive from reduced distractions in pre-training chunks. We report additional evidence for the burstiness property in Appendix D.
Note that duplication in pre-training sequences can also result in increased burstiness property, which may negatively impact the performance of language models. We analyse the distinct n-gram phrases of pre-training sequences in Appendix D and will investigate the impact of duplication using different pre-training corpora in future work.

heading: 6 Related Works
text: Instance-Level Pre-training Data Composition GPT-3 (Brown et al., 2020) was pre-trained by packed documents with causal masking, with the idea that not adopting any dynamic masking can improve pre-training efficiency. Current opensource pre-training frameworks, such as MegatronLM (Shoeybi et al., 2019), FAIRSEQ (Ott et al., 2019), EasyLM (Geng, 2023), LLM360 (Liu et al., 2023b), OLMo (Groeneveld et al., 2024) also follow this strategy for pre-training. In Levine et al. (2022), authors pair similar sentences within the
same sequence, while Gu et al. (2023) propose packing documents that contain similar intrinsic tasks for continual pre-training, improving the incontext learning ability of models. Recently, Shi et al. (2023) emphasise that packing relevant documents can enhance language models’ in-context learning and context utilisation; however, our findings indicate that packing documents can adversely affect performance, and learning each document independently using intra-document causal masking can reduce the distraction and improve the performance.
Distribution Properties of Pre-Training Data Chan et al. (2022) shows several data distribution properties can drive in-context learning ability, e.g., large numbers of long-tail classes, dynamic meanings of inputs, and Zipf’s distribution of class frequency. Han et al. (2023) used a gradient-guided method to select small-scale data for continual pretraining, showing data exhibiting burstiness properties can enhance in-context learning performance.
de Vries (2023) observes that over 80% of pretraining examples consist of fewer than 2K tokens, which suggests that the token distribution of individual documents is not well-suited for the longcontext window. To improve the long-context modelling ability of models, Staniszewski et al. (2023) retrieves related documents to construct fine-tuning data with long-range dependency.
Pre-Training Data Quality Gunasekar et al. (2023) selected high-quality data to pre-train a small-size coding model, achieving comparable performance with larger models. Shin et al. (2022); Gao et al. (2021) emphasised the importance of pre-training data diversity. Lee et al. (2022); Tirumala et al. (2023); Soboleva et al. (2023); Abbas et al. (2023) showed the importance of data deduplication on models’ generalisation. In our work, we use a diverse and high-quality pre-training dataset, namely SlimPajama (Soboleva et al., 2023), to highlight the importance of the sequence composition strategy on language model pre-training.

heading: 7 Conclusion
text: In this work, we investigate the impact of pretraining sequence composition by pre-training models from scratch. First, we find causal masking can result in unrelated documents distracting language modelling pre-training and hurting the performance on downstream tasks; we show that intra-document
8
causal masking can significantly improve the performance while decreasing the pre-training efficiency. Second, we find improving the relatedness of documents in pre-training chunks for causal masking pre-training can reduce some potential distractions in chunks; our proposed efficient retrievalbased packing method BM25Chunk can significantly improve the performance of language models without reducing pre-training efficiency.
Limitations
Efficiency of Intra-Document Causal Masking We show that intra-document causal masking is an effective method to improve the performance while decreasing the pre-training efficiency. We use FlashAttention2 (Dao, 2023) to implement intradocument causal masking masking without sacrificing too much efficiency (discussed in Appendix A). Still, we do not propose a method to solve this efficiency issue completely.
Objective of Sequences Construction. We discuss sequence construction methods, showing the importance of sequence compositions on the performance of models, but these methods lack an objective during sequence construction. Since specific data distribution properties may be related to models’ performance, we will explore using indicators of distributional properties to guide sequence construction in future works.
Scaling The Size of Language Models. Limited by the computation resources, we cannot conduct experiments on larger models with more pretraining steps, and different results might be drawn when increasing the models at a specific scale. However, this work could be directly valuable for investigating pre-training relatively small models that aim at facilitating the use of language models under resource-constrained conditions.

heading: Acknowledgements
text: PM was partially funded by ELIAI (The Edinburgh Laboratory for Integrated Artificial Intelligence), EPSRC (grant no. EP/W002876/1), an industry grant from Cisco, and a donation from Accenture LLP; and is grateful to NVIDIA for the GPU donations. YZ was supported in part by the UKRI Centre for Doctoral Training in Natural Language Processing, funded by UK Research and Innovation (grant EP/S022481/1) and the University of Edinburgh, School of Informatics. This work was
supported by the Edinburgh International Data Facility (EIDF) and the Data-Driven Innovation Programme at the University of Edinburgh. The authors extend their gratitude to Xiaomi AI Lab for their GPU donations and assistance; as well as to Tri Dao, Piotr Nawrot, Giwon Hong, Xiaotang Du, and Aryo Gema for their help and feedback.

heading: B Pre-Training Details
text: Hyperparameters In our experiments, we use the 1.3B model, which has 24 layers, a hidden size of 2048, and 16 attention heads. We use a batch size of 4 million tokens for both the models with 2K and 8K context window sizes and pre-train models using 150B tokens with 38400 steps, which costs 40 hours to pre-training a causal masking model using 128 80G A100 GPUs. We use Adam optimiser with β1 = 0.90, β2 = 0.95, a weight decay of 0.1, and a cosine learning rate scheduler. The peak learning rate is 3× 10−4, decreasing to 3× 10−5 at the end.
Pre-Training Corpus We sample documents with 150B tokens sampled from SlimPajama for pre-training. All models are pre-trained using the same set of documents. In Table 6, we present the number of documents and the token proportions for each subset.

heading: C Analysis of BM25Chunk
text: C.1 Time Complexity Analysis
In BM25, the similarity score between a query and a document is based on sparse representations, where each query and document is represented by the terms it contains; such sparse representations are stored in inverted indices, which map terms to the documents that contain them, along with necessary statistics such as the term frequency and the document frequency. The time complexity of computing similarities between a query and documents in BM25 using an inverted index is O(Q × K), where Q denotes the number of tokens in the query, and K represents the number of total documents.
To improve efficiency, we restrict BM25Chunk’s retrieval process within a document buffer rather than entire large-scale corpora. The buffer caches k documents, which enables similarity calculations between a term and documents to be at most k times. Since each query is a document, it could contain a large number of tokens; we remove the stop words and randomly sample q tokens to reduce the length. Therefore, the time complexity of sequence construction in BM25Chunk is reduced to O(q × k). In Figure 5, we test the sequence construction speed using different q and k.
C.2 Implementation Details
We randomly group documents in batches of 5000K and build indexes within each group. The BM25 indexes of pre-training corpora with 150B tokens require 244GB storage memory. For both 2K and 8K settings, the document buffer size k is 3072, and the maximum length of query q is 500. The
13
data construction speed is 50.0K tokens per second using 16 CPU cores, and speeds using different settings are presented in Figure 5.
C.3 Ablation Studies
Effectiveness of Document Buffer BM25Chunk conducts retrieval within a document buffer, which may result in retrieving less relevant documents, so we conduct experiments on different document buffer sizes to investigate its effectiveness. We conduct ablation experiments using 0.3B models with a context window of 2048, trained with 13B tokens, the compute-optimal number of tokens according to Hoffmann et al. (2022). We present the PPL improvement over UNIChunk on the validation set of SlimPajama in Table 7. The results show that retrieving from different sizes of document buffers can improve PPL, indicating the effectiveness of retrieving from a small-scale document set. BM25Chunk with a buffer size of 4096 achieves the best result, while increasing the size to 8192 does not improve the PPL.
Effectiveness of Retrieval BM25Chunk conducts multi-hop retrieval to retrieve a sequence of documents, which could potentially help models learn long-distance relationships across documents, and this benefit has been revealed by its high accuracy on HotpotQA, a multi-hop QA task, as shown in Section 4.3. An alternative choice is retrieving multiple documents at once to fill a pre-training chunk, and we present such one-hop retrieval in Table 7. The result indicates that BM25Chunk with multi-hop retrieval can improve the PPL more effectively. Besides, we experiment with random sampling documents from the buffers without retrieval; the result shows the effectiveness of retrieval.
Dense Retrieval Method An alternative retrieval method to BM25 is dense retrieval. We use Contreiver (Izacard et al., 2022) as the dense retriever and compare it with BM25. Following Shi et al. (2023), we embed pre-training documents to dense vectors using Contriever and use FAISS (Johnson et al., 2019) to accelerate the retrieval process instead of using the document buffer. Then, we construct pre-training chunks using the same process introduced in BM25Chunk. We present the result produced by the dense retrieval method in the last line of Table 7. We observe that the improvement of the dense retrieval method is less than BM25.

heading: D Analysis of Data Distribution Properties
text: Chunk Frequency In addition to Zipf’s coefficient, we analyse the burstiness property through the chunk frequency of tokens. Specifically, chunk
14
frequency refers to the number of chunks where a specific token appears. Given a corpus, if a specific token appears in fewer chunks, it indicates more concentrated occurrences in chunks containing the token, demonstrating a higher burstiness property. In Figure 6, we can see that low-frequency tokens appear in fewer chunks in BM25Chunk compared to MIXChunk and UNIChunk, indicating these lowfrequency tokens are gathered through the retrievalbased construction method.
Distinct N-gram The burstiness property can correlate to the duplication in a sequence, which may negatively affect models, e.g., models may tend to copy phrases from context. We use SlimPajama, a high-quality and deduplicated dataset, as the pretraining corpus, which can alleviate the duplication issue in BM25Chunk. We use the percentage of distinct n-grams within a sequence to analyse the duplication issue, as shown in Table 10. The results show that, with BM25Chunk, pre-training sequences contain a lower percentage of distinct n-grams than MIXChunk and UNIChunk.

heading: E Analysis of Distraction Proportions in Different Settings
text: In Figure 7, we report the average distraction proportion (defined in Equation (2)) over layers us-
ing different settings. Specifically, we analyse distraction proportions in different settings by varying the 1) modalities of corpus: text and code using Wikipedia and GitHub; 2) the separator token: [EOS] and line break token \n; 3) the length of distraction context, |Cd| = 256 and 512.
Comparing different separators [EOS] and \n, (a) (e), (b) (f), (c) (g), and (d) (h), we observe that causal masking models can obtain lower distraction proportions using [EOS], indicating causal masking models can benefit from [EOS] to ignore irrelevant context during pre-training. We present the impact of changing the separator from [EOS] to \n on PPL and distraction proportion in Table 8. The results show that PPL and DISTPROP increase after the replacement for causal masking models, while INTRADoc obtains better results using \n as the separator since it does not train on sequences
15
where documents are separated by [EOS] using intra-document causal masking.
Comparing Wikipedia (a) (b) (e) (f) and GitHub (c) (d) (g) (h), MIXChunk is more distracted by the irrelevant context in code generation.
Comparing different length distraction contexts, (a) (b), (c) (d), (e) (f) and (g) (h), models are more distracted when |Cd| increases, while much better than the baseline of uniform distribution y = |Cd|/(|Cd|+ x).
Comparing INTRADoc (red line) and causal masking models, we observe that intra-document causal masking results in significantly lower distraction proportions in all cases. This phenomenon may imply that using causal masking without considering the boundaries of documents negatively impacts language modelling performance, and the models can be more robust to irrelevant contexts when increasing the relatedness of documents in pre-training chunks.

heading: F Next Token Accuracy of Pre-Trained Language Models
text: In addition to PPL, we report the next token accuracy of pre-trained language models in Table 9.
16

heading: Future Work
text: Since specific data distribution properties may be related to models’ performance, we will explore using indicators of distributional properties to guide sequence construction in future works. Scaling The Size of Language Models. Limited by the computation resources, we cannot conduct experiments on larger models with more pretraining steps, and different results might be drawn when increasing the models at a specific scale. However, this work could be directly valuable for investigating pre-training relatively small models that aim at facilitating the use of language models under resource-constrained conditions