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

Abstract: The prohibitive training costs of 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 LLM 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: https://github.com/shaochenze/PatchTrain.

• The paper introduces patch-level training that compresses tokens into patches, reducing computational costs by nearly 50% while preserving model performance.
• The two-stage method first optimizes patch-level prediction then refines parameters via token-level training for effective knowledge transfer.
• Experimental results highlight enhanced neuron activation and swift scaling with more data, demonstrating the method's efficiency and practical impact.

Patch-Level Training for LLMs: An Efficient Approach

The paper Patch-Level Training for LLMs by Chenze Shao et al. tackles a critical concern in the field of LLMs - the computational cost of training. This paper presents an innovative training strategy termed "patch-level training," aimed at enhancing training efficiency without compromising model performance. The core idea revolves around compressing multiple tokens into patches for initial training, significantly reducing computational loads by processing shorter sequences.

Methodology

Patch-Level Training Mechanism

The conventional token-level training of LLMs requires each token to be processed individually, leading to considerable computational expenses due to the extensive token sequences involved. Shao et al. propose an alternative training paradigm where sequences of tokens are compressed into patches. This patch-level training involves constructing shorter sequences of patches and training the model to predict the subsequent patch.

This training method is divided into two stages:

1. Patch-Level Training: The model is fed shorter sequences of patches, compressing groups of $K$ tokens into single patches, and it learns to predict the next patch.
2. Token-Level Training: Subsequently, a smaller subset of data is used to continue with token-level training. This step aligns with the model's inference mode, refining the parameters initialized during the patch-level training to adapt to token-level inputs.

The trained parameters from the patch-level model (PLM) are harnessed to initialize the token-level model (TLM), ensuring a seamless knowledge transfer.

Experimental Validation

Training Efficiency

The substantial reduction in computational requirements through patch-level training is illustrated with LLMs of varying sizes (ranging from 370M to 2.7B parameters). By setting the patch size $K=4$ and using this method on $\lambda=2/3$ of the training data, the paper demonstrates that the computational cost can be halved without any significant detriment to performance. Empirically, this approach maintained comparable perplexity scores and even enhanced model performance on zero-shot NLP benchmarks.

Performance Across Different Settings

The experiments underscore that patch-level training manifests favorable scaling properties, particularly when ample training data is available. Notably, the performance benefits from an increase in data size more swiftly than traditional token-level training. However, the model performance gain due to an increase in model size was relatively modest in comparison.

Detailed Analysis

Hyperparameters: Patch Size and Data Fraction

1. Patch Size ($K$): The findings suggest that a patch size of $K=4$ achieves a balance between efficiency and performance. Larger patch sizes (e.g., $K=16$) exhibit slight performance degradation, possibly due to the increased challenge in transferring knowledge to the token-level model.
2. Data Fraction ($\lambda$): The hyperparameter $\lambda$ determines the proportion of patch-level to token-level training data. An optimal $\lambda$ of around 2/3 was identified under fixed computational budgets, providing a beneficial trade-off between training cost and model performance.

Neuron Activation Perspective

To elucidate why patch-level training yields better learning efficiency, the paper presents an analysis of neuron activation. Training with patches leads to increased neuron activation rates, particularly in earlier transformer layers, indicating a denser and potentially more effective utilization of model parameters during the knowledge acquisition phase.

Implications and Future Directions

Patch-level training heralds a significant advancement in the efficiency of training large-scale LLMs. Practically, this approach can halve training costs without sacrificing model accuracy or performance, thereby making LLM development more accessible and sustainable. Theoretically, it opens avenues to explore neuron activation patterns and their impacts on efficient learning.

Future research could focus on validating the scalability of patch-level training with model and dataset sizes akin to state-of-the-art LLMs. Understanding the empirical scaling laws that incorporate varying $K$ and $\lambda$ values could further optimize training costs. Additionally, investigating the applicability of patch-level training to other domains, such as image or speech processing, would expand the horizons of this method's utility.

In conclusion, the paper by Shao et al. presents a substantial contribution to efficient LLM training methodologies. By innovatively restructuring the training process with patch-level sequences, they achieve a balance of reduced computational burdens and maintained model performance, laying the groundwork for future efficiency-focused advancements in machine learning.