Structured Packing in LLM Training Improves Long Context Utilization

We thank the Reviewer for a thoughtful review.

We are not sure what point 1 of weaknesses means. We gently ask the reviewer for clarification.

We agree that perplexity is not enough to fully evaluate the method. We now provide positive results for two datasets: TREC [1, 2] and Qasper [3, 4]. See details in the general answer and the revised pdf. In the general answer, we also provide new experiments with neural retrieval, which allow for efficient scaling (see the table below).

We consider the techniques relying on extensive conversations (e.g. LongChat [5]) as complementary to our work. We note that SPLiCe can utilize data that already exists in vast quantities and can be easily applied to different types of text (like code, Wikipedia articles, stack exchange questions and answers, etc.). It is interesting to evaluate how both would interact, which we now include in the future work section on page 9.

Answers to the questions:

Q1: We now test our method using Contriever-MSMARCO [6] instead of BM25 and observe that the results are very close. We attach the details in the general response and in the revised version of the paper (see Table 1 on page 5).

Q2: SPLiCe using the Contriever-MSMARCO retrieval is scalable to large datasets. Namely, the dataset can be effectively queried using an approximate index (Faiss [7] in our case) which is a common technique, used in e.g. kNN-LM [8], RETRO [9]. We have also performed an experiment to check whether SPLiCe will bring improvements when applied to a random subset of a larger set. For this, we have extracted 2GB from English Wikipedia and augmented it using SPLiCe. The results are presented in Table 3 on page 6. The outcome indicates that SPLiCe does not need to retrieve from the whole dataset to provide performance improvements. For example, one could split the training dataset into smaller parts and run SPLiCe independently for each of them reducing the complexity even further.

[1] Li et. al. Learning Question Classifiers
[2] Hovy et. al. Toward Semantics-Based Answer Pinpointing
[3] Dasigi et. al. A Dataset of Information-Seeking Questions and Answers Anchored in Research Papers
[4] Shaham et. al. SCROLLS: Standardized CompaRison Over Long Language Sequences
[5] Li et. al. How Long Can Open-Source LLMs Truly Promise on Context Length?
[6] Izacard et.al. Unsupervised Dense Information Retrieval with Contrastive Learning
[7] Johnson et. al. Billion-scale similarity search with GPUs
[8] Khandelwal et. al. Generalization through Memorization: Nearest Neighbor Language Models
[9] Borgeaud et. al. Improving language models by retrieving from trillions of tokens

We hope that our explanations and additional results show the potential of SPLiCe. In particular, our evaluation on additional downstream tasks along with results regarding Contriever-MSMARCO and Faiss replacing BM25.

Considering the approaching rebuttal deadline, we ask the Reviewer for comments and a score increase in case our findings meet the requirements.

We thank the Reviewer for their insightful feedback.

We note that SPLiCe is a stand-alone solution. Namely, it could be used for the whole training. We use it in fine-tuning due to computational constraints. We also note that it is fairly standard to train on a short context for the majority of training steps and extend it in the final phase, see e.g. CodeLlama [1]. We believe that the performance gains validate the usage of the algorithm, taking into account that it is simple and its computational overhead is negligible compared to the training costs.

We totally agree with the request for NLP evaluation. We now provide results for TREC [2, 3] and Qasper [4, 5] datasets, which confirm the usefulness of the SPLiCe training. Please refer to the results in the general answer.

In what follows, we address the Reviewer’s questions:

• Regarding the imbalance. We use each document only once. A document that is used in training is removed (see the last lines of Alg 1). This design choice ensures that for one epoch the whole dataset is seen. What is more, we mix the SPLiCe prepared data with data prepared in a standard way, which should prevent the model from overfitting to the structure of SPLiCe constructed examples. In our setup, we have not noted training instabilities. We agree that there is some subtlety here and open questions, which we now note in the further research section (see page 9).
• We agree that there is some risk in using ‘unnatural’ packed documents. To alleviate this we train on a mixture of data prepared in the standard way and SPLiCe (which is now explained in Section 3.1). Based on [1, 6] we think that using data mixture is enough to alleviate false biases.
• SPLiCe is a general framework for constructing long context data using short documents. For example, for k=1 and the retriever returning elements from the closest subdirectory in the repository we get the Repo method [7]. With larger k and reversed order we get the data that resembles inputs to retrieval augmented models (e.g. REALM [8]). We believe that leaving the possibility for other choices paves the way for future research and is of use to the community.
• Thank you for pointing out several exposition improvement areas. We have made multiple changes towards this goal (e.g. the mentioned split of the results in Table 1 and Table 2 and the titles) and we will gladly accept any further suggestions.

[1] Rozière et. al. Code Llama: Open Foundation Models for Code
[2] Li et. al. Learning Question Classifiers
[3] Hovy et. al. Toward Semantics-Based Answer Pinpointing
[4] Dasigi et. al. A Dataset of Information-Seeking Questions and Answers Anchored in Research Papers
[5] Shaham et. al. SCROLLS: Standardized CompaRison Over Long Language Sequences
[6] Ouyang et. al. ”Training language models to follow instructions with human feedback
[7] Wu et. al. Memorizing Transformers
[8] Guu et. al. REALM: Retrieval-Augmented Language Model Pre-Training

Using the additional time, we managed to evaluate our method using larger 7B parameter models, longer context and more downstream tasks (the details are present in general response). We hope that our additional results and explanations showcase the potential of SPLiCe.

Considering the approaching rebuttal deadline, we ask the Reviewer for comments and a score increase in case our findings meet the requirements.

We thank the Reviewer for an encouraging review.

We now provide results showing that the SPLiCe training is useful for downstream NLP tasks: TREC [1, 2] and Qasper [3, 4]. Likewise, we are happy to report results for SPLiCe using Contriever-MSMARCO [5]. Please see the details in the general answer and the revised pdf and let us know if this is satisfactory.

Due to computational constraints, we have not been able yet to go beyond the current size/context length but we hope to provide results for the 7B/32K model and 270M/64K model before the rebuttal deadline.

Regarding questions:

Q1: We acknowledge that this statement is confusing. We have rephrased it as follows: "Despite the non-code nature of StackExchange we still note perplexity improvements on code. Unsurprisingly, the perplexity on code is not as good as when training directly on code data.”.

Q2: We apologize, but a mistake has slipped through, and we have unintentionally provided results of the 3B parameter model that was trained with a context length smaller than 32K. We corrected that in the revised version, providing the corresponding perplexity improvement results that now match the description. We also re-evaluated the model on the kv retrieval task and added the comment about improved performance. See the revised pdf, page 8 and 20.

Q3: We thank you for spotting the typo. We have fixed it and merged sections 3.2 and 3.3 into one.

Q4: Thank you for the suggestion. We have rephrased this part as follows: “Our results indicate that SPLiCe enhances model performance and can be used to train large models to utilize long contexts better.”

[1] Li et. al. Learning Question Classifiers
[2] Hovy et. al. Toward Semantics-Based Answer Pinpointing
[3] Dasigi et. al. A Dataset of Information-Seeking Questions and Answers Anchored in Research Papers
[4] Shaham et. al. SCROLLS: Standardized CompaRison Over Long Language Sequences
[5] Izacard et.al. Unsupervised Dense Information Retrieval with Contrastive Learning

Using additional time, we extended our research to 7B parameter models with 32K context and 270M parameter models with 64K context length. We add additional results to the general response. We hope that those findings along with the ones presented in our previous responses showcase the potential of SPLiCe. In particular our previous results have tested the method on two additional downstream tasks.

Considering the approaching rebuttal deadline, we ask the Reviewer for comment and a score increase in case our findings meet the requirements.

Table: Training Methods and Perplexity Results for Various Programming Languages

We train a 270M parameter model on a 50/50 mixture of the RedPajama data (organized in a standardway) and code data (organized in one of four ways: SPLiCe BM25, SPLiCe Contriever-MSMARCO, Repo, Baseline – standard). We evaluate the perplexity of the models on arXiv and various subsets of StarCoder. The boldface denotes the best result within the training mixture. Note that SPLiCe is generally on par or better than Repo and beats the Baseline by a large margin. For detailed results, see Appendix B.

7B parameter model

We have managed to tune the 7B parameter OpenLLaMA v2 model for 32K context length using 2B tokens 50/25/25 mixture of Redpajama prepared in standard way and StackExchange, C - prepared using SPLiCe. We compared the performance on TREC of the SPLiCe model to the model that was trained on data prepared in a standard way. The results are available in the table below. For each context length, we average the results across $10$ sets of in-context examples.

64K context length

We have managed to tune 270M parameter models for 64K context following the approach presented in the paper. The evaluation results are available in the table below.

We plan to add those additional results to the pdf before the rebuttal deadline.

Due to the approaching rebuttal deadline, we have uploaded the PDF version without color-marked differences. We have also added the 7B parameter model results and results with the 64K context for 270M parameter models. Those results were previously stated in our comments.