Beyond Next Token Prediction: Patch-Level Training for Large Language Models

Authors: Chenze Shao, Fandong Meng, Jie Zhou
Primary area: foundation or frontier models, including LLMs
Keywords: large language models, patch-level training

Abstract:
The prohibitive training costs of Large Language Models (LLMs) have emerged as a significant bottleneck in the development of next-generation LLMs. In this paper, we show that it is possible to significantly reduce the training costs of LLMs without sacrificing their performance. Specifically, we introduce patch-level training for LLMs, in which multiple tokens are aggregated into a unit of higher information density, referred to as a `patch', to serve as the fundamental text unit for training LLMs. During patch-level training, we feed the language model shorter sequences of patches and train it to predict the next patch, thereby processing the majority of the training data at a significantly reduced cost. Following this, the model continues token-level training on the remaining training data to align with the inference mode. Experiments on a diverse range of models (370M-2.7B parameters) demonstrate that patch-level training can reduce the overall training costs to 0.5$\times$, without compromising the model performance compared to token-level training. Source code: \url{https://github.com/shaochenze/PatchTrain}.

Decision: accept (spotlight)
Reviewer scores: 6, 8, 8, 8

Reviews and discussion:
<thread>
creation date:
1730390671787

thread_type:
official_review

<thread_object>
creator: reviewer_ktbd

<content>
summary:
This paper proposes a strategy to reduce training cost of LLMs using patch level training, where 'patch' refers to a sequence of consecutive tokens. The key idea is to perform training at the patch level initially whereby prior patches are used to predict the next patch, and then proceed with regular token level training.  The paper applies the approach to a range of model sizes (370M-2.5B) and shows that it can reduce the training cost by 50% without any performance degradation on model perplexity or downstream task performance. The paper presents variation in performance wrt patch size, model size, amount of training data and multi-epoch scenarios.

strengths:
* Presents a simple approach to reduce training cost using a patch level training for initialization the regular token-level training
* Reports experiments on a range of model sizes showing that the approach can reduce training cost by 50% without any degradation in model performance

weaknesses:
* There are some parts of the paper which are unclear. See questions below.
* It would be good to see some additional analysis/ablations to understand why the approach can lead to performance improvements.

questions:
* L156: "Instead, we maintain a single output head and make its prediction cover all tokens in the next patch". How is this done? Does your output vocabulary have a size of K*|V| where |V| is the number of symbols ?
* Equation 3 uses 'p' to refer to both probabilities and patches. Consider using different symbols for the two.
* Sec 3.7: Have you run experiments adding more transformer layers for mapping the token embeddings to patch embeddings and vice-versa at the output side? It is possible that a linear mapping is not powerful enough.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730411501836

thread_type:
official_review

<thread_object>
creator: reviewer_o6fe

<content>
summary:
The paper presents an interesting approach to accelerate the costly pre-training of language models. The core idea is to change the basic unit from a token to patch which comprises several tokens. This results in computational savings for the main transformer model (the input encoding and the model prediction still produces tokens but it does so based on a patch encoding).

strengths:
The paper presents an interesting approach to accelerate language model training, specifically:
* the approach appears to be novel
* the paper is well written and clear in the presentation
* the approach is thoroughly evaluated (different (small) model sizes, including an array of ablations, an alternative approach to computing patches etc.)

weaknesses:
* The evaluation is on the smaller scale of language mode sizes. Evaluation for larger model sizes is very challenging but it would be a large plus to understand the scalability of the method.
* Table 3/4 show PPL results and draw conclusions based on the difference between PPLs in the range of ~10 and later on in ~7. This assumes that PPL is a linear measure but it is in fact an exponential quantity based on cross entropy. Therefore, the actual entropy difference between PPL 11 and PPL 10 is far lower than it is between, say PPL 8 and PPL 7. I would encourage the authors to simply report entropy instead of PPL for this particular analysis to make this clearer.

questions:
* Line 266 mentions that the context length in the patch level training has little effect. Using context length KT performed as well as just T. However, these conclusions were drawn based on a very small model with 370M parameters which may not be able to exploit the larger context as well as higher capacity models. Did you perform a comparison on a larger model?
* Does the patch level training objective have any notion of which token position is currently being predicted within a patch (eq 3)? It appears to treat a patch as a bag of words with no order? This would be interesting to contrast with a setting where there is a notion of order.
* Section 3.5 analyses the patch size K and deems K=2/4 to have a good speed/accuracy tradeoff in your setting which is a two-stage procedure: (1) train with some K>1 and then (2) use K=1 at the end of training. However, you could also consider a more general setting where there are many more stages and where where you start with very large K and then progressively decrease K until you end up with K=1. Have you considered this?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730584305487

thread_type:
official_review

<thread_object>
creator: reviewer_zf1q

<content>
summary:
This work introduces a novel pretraining scheme aimed at reducing computational cost of pretraining. The core idea is to reduce the effective length of training sequences by concatenating and averaging the token representation of mutliple tokens into a patch representation. Then the model predicts tokens from the next patch using the output representation of the next patch. The simplicity of this idea is very intriguing as it does not require any extra parameters in the model. After such path-level pretraining with higher efficiency authors go back to tokenlevel training in order to ensure that the model will adjust to next token prediction representation.

Authors completed experiments showing that their method scales well with larger models and supports both streaming and multi-epoch training scenario. Extensive benchmarking showed that models trained with their approach are cheaper to train and are better or on the same level of performance as baseline models.

strengths:
1. the method is easy to implement and brings original contribution, i see it being used a lot in the community and industry if it will pass the proof of concept in real world scenario.
2. High quality experiments cover wide range of model sizes and data settings. In addition, wide range of benchmarks are covered to evaluate the final model performance compared with usual token level training.
3. Assuming that this approach would scale with even larger models, this work provides a significant contribution in the pretraining research.

weaknesses:
* Unclear conclusions about the effect of architecture. In the architecture ablation the observations are suggesting that there are might be some use cases where subsequent token level training does not work very well. IMO this might be important to highlight better in the text to avoid overselling the method.
* No numbers showing the actual speed improvements. Since all token projections are computed in patch level training, the expensive softmax operations over vocabularies have to be done. So it is interesting what is the speed improvement factor as we change K (patch size) in practice on real hardware.

questions:
* Following my last weakness point, did you measure the speed up coming from patch level training in practice? e.g. in tokens per second metric or so. This would be very interesting to know.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730754561674

thread_type:
official_review

<thread_object>
creator: reviewer_qqzw

<content>
summary:
This paper seeks to reduce the training cost of LLMs by using patch-level training in the initial stage. A patch consists of multiple consecutive tokens, and patch-level representation can reduce training costs due to shorter sequence representations. Experiments performed across four model sizes and six NLP tasks show that the proposed method can reduce training costs by up to 50% with little to no performance degradation.

strengths:
1. Reducing LLM training costs is an important research question. The problem formulation, motivation, and improvement objective of the paper are clear.

2. The proposed method is described clearly, and there seem to be adequate details to reproduce the work. The cost reduction of the proposed training method is significant and robust across multiple LLM sizes and NLP tasks.

3. The experiments cover a wide range of settings such as the fraction of patch-level training (lambda), patch size (K), model size, etc.

weaknesses:
1. The optimal hyperparameters and intuitions derived from the work may be very specific to the settings associated with the paper. For example, the optimal values of lambda and patch size may be sensitive to the tokenizer. Therefore, the improvements of the proposed work may not generalize similarly to a new tokenizer.

2. It is not obvious whether the patch-level pretraining is suitable for larger scales. There are issues with training stability and convergence that only occur when the model size is scaled to several billions.

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1732015409913

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
We sincerely thank all the reviewers for their great efforts in reviewing our work. In response to the valuable feedback, we have made several major updates to our manuscript, as outlined below:

1. In Appendix A, we provide the pseudocode for patch-level training to enhance its clarity.
2. In Appendix B, we give a detailed analysis of the practical speed improvements through patch-level training.
3. In Appendix C, we introduce an alternative mixup training strategy that we have recently explored. This section also includes a comparative analysis of its performance against patch-level training.
4. In Appendix D, we analyze the effect of tokenizer on patch-level training.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732162457378

thread_type:
official_comment

<thread_object>
creator: area_chair_zkdq

<content>
title:
Please respond and update the score if necessary

comment:
Dear Reviewers,

Kindly ensure that you respond proactively to the authors' replies so we can foster a productive discussion. If necessary, please update your score accordingly. We greatly appreciate the time and effort you’ve dedicated to the review process, and your contributions are key to making this process run smoothly.

Thank you,

AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces a pre-training method designed to reduce the training costs of large language models by implementing patch-level training, where multiple consecutive tokens are aggregated into a single patch. This approach can decrease computational expenses by up to 50% while maintaining performance levels. Experiments show the method scales effectively with larger models and supports both streaming and multi-epoch training scenarios, providing cost-effective solutions without additional model parameters. The models trained using this technique perform as well as or better than traditional baseline models.

The paper's method effectively reduces LLM training costs without compromising performance. It is thoroughly evaluated across a range of model sizes, settings, and benchmarks, showcasing robust performance. The approach is straightforward to implement and could achieve widespread adoption in both research and industry, making a substantial contribution to the field of pretraining research. The paper could benefit from additional analysis involving larger models.

Additionally, the authors have addressed most of the reviewers' questions thoughtfully and thoroughly. I recommend accepting this paper, and I am excited to see it presented at the conference.