Paper: Structured Packing in LLM Training Improves Long Context Utilization (ICLR-2024)

Reviewer comment (Novelty):
Using Lexical matching methods to concatenate the documents into a longer one is a very engineering technique and it is a straightforward solution to construct longer samples.

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

Reviewer comment (Experiments):
The experimental results are almost based on PPL, lacking experiments on real-world tasks. More experiments on benchmarks such as zeroScrolls[1] or L-Eval[2] to validate their models are needed. More extensive testing across a broader range of tasks and datasets would provide a more comprehensive evaluation of the method.

Author response:
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).

Reviewer comment (Experiments):
Presently, the prevalent strategies for training long context models involve the use of extensive conversations and literary works. A comparative analysis of SPLICE with these existing methodologies is thus a necessary step.

Author response:
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.

Reviewer comment (Experiments):
How can we SPLICE on very large pertaining corpus which usually has more than 400B tokens?

Author response:
**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.

Reviewer comment (Evaluation):
Considering the additionally introduced complexity (though the SPLiCE algorithm is simple), the performance improvement looks very marginal, especially for large-scale models.

Author response:
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.

Reviewer comment (Experiments):
Only the part of the training is replaced with SPLiCE from the random baseline. That might be one of the reasons for marginal improvement, but it also implies that SpLiCE is not a standalone solution that can completely replace the existing training algorithm.

Author response:
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.

Reviewer comment (Evaluation):
Language modeling perplexity is the main evaluation metric. Comparing performance on other NLP downstream tasks that require long context modeling might be better to evaluate the effectiveness of the proposed method.

Author response:
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.

Reviewer comment (Experiments):
I raised several concerns about why SPLiCE is not sufficient (or at least not fully validated) as it is. Could you address them?

Author response:
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.

Reviewer comment (Experiments):
I guess the number of neighbors for each document is skewed, meaning that there exists hub documents. In that case, although a root document is randomly sampled from the document corpus, the retrieved documents are not uniformly distributed in terms of their likelihood and order. Couldn't this be a problem that may result in an imbalance in training?

Author response:
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).

Reviewer comment (Experiments):
Packed documents are unnatural and different from contiguous documents. Is there any way to alleviate this issue?

Author response:
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.

Reviewer comment (Experiments):
As expected, using related documents in a long context is better than the random baseline. However, any design choices (top-k, order, or even REPO vs. SPLiCE) give clear differences. In particular, top-1 is the best, and in that case, BFS is the same as DFS.

Author response:
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.

Reviewer comment (Presentation):
Why is Table 1 required? Table 2 fully covers Table 1.

Author response:
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.

Reviewer comment (Presentation):
The structure of the paper can be improved. For example, it is awkward that Section 4 also includes experiments while the title of Section 3 is experiments. Also, multiple Figures and Tables can be merged to spare some space for more extensive experiments or discussion.

Author response:
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.

Reviewer comment (Experiments):
The effectiveness of the method is only validated on language modeling and the key-value retrieval task. This does not guarantee the resultant model is stronger on realistic use cases. To test the usefulness of SPLICE, I would highly recommend comparing the models on more challenging and realistic long-context downstream tasks such as Quality and Squality.

Author response:
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.

Reviewer comment (Experiments):
It would be great if the method is tested on more settings: Use a neural retriever in addition to BM25, go beyond 3B and 32K, etc.

Author response:
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.

Reviewer comment (Writing):
Typo: 3.2 “Moreover, even thorough this method uses only the code data”

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

Reviewer comment (Writing):
In the abstract “Our results indicate that SPLICE enhances model performance across various task” - The context of this sentence is the 270M model. It is in fact only tested on one task: language modeling (though there are different datasets). You might want to rephrase to reduce confusion.

Author response:
**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.”