Patch-Level Training for Large Language Models

## Abstract

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: https://github.com/shaochenze/PatchTrain .

## 1 Introduction

Large Language Models (LLMs, Achiam et al., ; Touvron et al., ; ; Team et al., ; Bai et al., ) have achieved remarkable progress in language understanding and generation, which are primarily attributed to their unprecedented model capacity and the corresponding growth in the volume of training data they require (Kaplan et al., ; Hoffmann et al., ). However, this scaling up comes with a substantial rise in computational costs, making the training efficiency of LLMs a critical concern. Despite the ongoing efforts on efficient LLMs (Wan et al., ), it remains a formidable challenge to reduce training costs without compromising the model performance.

Specifically, the amount of compute (FLOPs) required for training LLMs is approximately proportional to both the number of model parameters $N$ and the number of text units (i.e., tokens) $D$ in the training data. This relationship can be expressed as:

$$ $C\approx 6ND.$ (1) $$

Therefore, strategies for reducing training costs can target either the reduction in the number of model parameters $N$ or the number of text units $D$. One prominent approach to reduce the parameter size $N$ is called model growth (Gong et al., ; Yang et al., ; Chen et al., ). Rather than equipping the model with full parameters from the beginning, it advocates for a progressive expansion of the model’s parameter size throughout the training phase, thereby reducing the average parameter size during training. Nonetheless, a model’s performance hinges on an adequate number of parameters to store extensive knowledge and develop intricate reasoning capabilities. The inadequacy of parameters inherently limits the scope of knowledge and capabilities a model can develop, rendering it challenging to match the performance of training with the full parameter set.

The second pathway to lowering training costs is reducing the number of text units $D$ within the training data. This direction remains largely unexplored but intuitively holds more promise, as the knowledge embedded within training data is sparsely distributed across numerous tokens, with each token encapsulating a minimal amount of information. This sparse distribution of information results in a scenario where, despite the substantial computational costs incurred during each learning step, only a small fraction of model parameters that are relevant to the current token are effectively updated. By increasing the amount of information the model processes at each learning step—that is, by augmenting the information density of text units and thus reducing the number of text units $D$—we could potentially boost training efficiency significantly without sacrificing model performance.

Building on these insights, this paper introduces patch-level training for large language models, 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. Specifically, we divide the training process into two stages: patch-level training and token-level training. 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. The resulting parameters are used to initialize the token-level model, which then continues training on the remaining data to adapt the knowledge gained during patch-level training to the token-level.

Figure 1 illustrates the efficiency advantage of patch-level training, where the area of the shape represents the overall training costs. With a patch size of $K$, the amount of compute required for patch-level training is $1/K$ of that required for token-level training. When a fraction $\lambda$ of the training data is compressed into patches, the overall training costs are reduced to $\lambda/K+1-\lambda$ times the original costs. For instance, to halve the training costs, one could set the patch size $K=4$ and conduct patch-level training on $\lambda=2/3$ of the training data.

Figure: Figure 1: Visualization of overall training costs with patch compression for a fraction $\lambda$ of training data and patch size $K$.
Refer to caption: extracted/5851977/intro.png

Figure: Figure 2: Negative log-likelihood (NLL) loss on test set w.r.t the number of processed tokens during the training of 370M-parameter Transformers.
Refer to caption: extracted/5851977/loss.png

Employing the above settings ($K=4,\lambda=2/3$), we train a series of LLMs of varying sizes (370M-2.7B parameters) on the Pile dataset (Gao et al., ). Figure 2 illustrates the trend of NLL loss against the number of training tokens for the 370M model. After initialization with patch-level training, the model experiences a rapid decrease in loss as it continues token-level training on the remaining data. Remarkably, it achieves an even lower loss in comparison with training from scratch, while reducing training costs by half. By further adjusting the hyperparameter settings, even higher acceleration rates can be achieved, with only a slight sacrifice in model performance.

## 2 Patch-Level Training

In this section, we outline the patch-level training approach for large language models, as illustrated in Figure 3. Initially, the token sequence is transformed into a patch sequence by compressing every $K$ consecutive tokens into a single patch. This patch sequence is then fed into the sequence model, and the model is trained to predict all tokens in the next patch. The knowledge acquired during patch-level training is subsequently transferred to the token-level model. Specifically, we use the parameters obtained from the patch-level model to initialize the token-level model, and then proceed with token-level training on the remaining data.

Figure: Figure 3: Overview of patch-level training. Every consecutive $K$ token embeddings are averaged to form the patch embedding. The sequence model is fed the patch sequence and trained to predict the next patch. The cross-entropy loss is computed based on each patch prediction vector and all the subsequent $K$ tokens in its next patch.
Refer to caption: extracted/5851977/model.png

While formulating the patch-level model structure, our goal is to minimize the discrepancy between patch-level and token-level models, thereby ensuring that the knowledge gained during patch-level training can be smoothly transferred to the token-level model. Given the context length $T$ for token-level training, we set the context length for patch-level training as $KT$, which is then compressed to a patch sequence of length $T$ to maintain consistency with the subsequent token-level training. To avoid introducing unnecessary parameters during token-to-patch compression, we represent the patch embedding as the average of its associated token embeddings. Let $p_{i}$ be the $i$-th patch, $x_{iK+k}$ be the $k$-th token in the $i$-th patch, and $E$ be the embedding function. The patch embedding is:

$$ $E(p_{i})=\frac{1}{K}\sum_{k=0}^{K-1}E(x_{iK+k}).$ (2) $$

The patch-level model is trained through next patch prediction, i.e., predicting the $K$ tokens in the next patch. The simultaneous prediction of multiple tokens has been explored in speculative decoding, which typically employs multiple output heads and each head is responsible for predicting a distinct token (Cai et al., ; Lin et al., ). However, this approach would also entail additional parameters that may be unfavorable for the subsequent knowledge transfer. Instead, we maintain a single output head and make its prediction cover all tokens in the next patch. Specifically, we calculate the cross-entropy loss for all the subsequent $K$ tokens based on the same patch prediction $p_{\theta}(\cdot|p_{<i})$, resulting in the following loss function:

$$ $\mathcal{L}_{patch}=-\sum_{i=1}^{T}\sum_{k=0}^{K-1}\log p_{\theta}(x_{iK+k}|p_ {<i}).$ (3) $$

Since the model finally works at the token-level, it is essential to reserve some training data to adapt the patch-level model to token-level. Specifically, we conduct patch-level training on a fraction $\lambda$ of the training data, and then use the resulting parameters to initialize the token-level model. Following this, the token-level model continues training on the remaining data to adapt the knowledge gained during patch-level training to the token-level. As illustrated in Figure 1, the overall training costs are reduced to $\lambda/K+1-\lambda$ times the original costs of token-level training. When the amount of training data is limited, this approach can also be utilized for efficient multi-epoch training. For example, given a budget of $N$ epochs, we can conduct patch-level training on the first $N\lambda$ epochs, and then switch to token-level training for $N(1-\lambda)$ epochs.

## 3 Experiments

### 3.1 Setup

Datasets. We evaluate our approach on standard language modeling tasks, using the Pile dataset (Gao et al., ) containing 360B tokens for training (^1^11Previous works generally refer to the Pile dataset as having 300B tokens, but our actual measurement is 360B. The discrepancy is likely due to differences in tokenizers; we use the LLaMA2 tokenizer, which has a relatively small vocabulary, possibly resulting in more tokens. The perplexity scores are also incomparable with models using other tokenizers.). We assess the performance of LLMs from multiple aspects, including their perplexity, zero-shot accuracy, and instruction-following ability. Perplexity is calculated on the WikiText-103 test set (Merity et al., ).
We evaluate the zero-shot capabilities of language models on 6 NLP benchmarks, including MMLU (Hendrycks et al., ), HellaSwag (Zellers et al., ), PIQA (Bisk et al., ), WinoGrande (Sakaguchi et al., ), ARC-E, and ARC-C (Clark et al., ) (^2^22https://github.com/EleutherAI/lm-evaluation-harness). For the pre-trained LLMs, we conduct instruction fine-tuning using the Alpaca dataset by GPT4 (Taori et al., ), and then evaluate their instruction-following abilities on MT-Bench (Zheng et al., ).

Models. We use the Transformer backbone (Vaswani et al., ) and adopt most of the architecture designs from LLaMA (Touvron et al., ). We apply pre-normalization using RMSNorm (Zhang & Sennrich, ), use the SwiGLU activation function (Shazeer, ), and rotary positional embeddings (Su et al., ). We also apply FlashAttention-2 (Dao, ) to accelerate attention computation. We scale the model demension and obtain 4 different sizes of Transformers: Transformer-370M (hidden_size=1024, intermediate_size=2752, hidden_layers=24, attention_heads=16), Transformer-780M (hidden_size=1536, intermediate_size=4128, hidden_layers=24, attention_heads=16), Transformer-1.3B (hidden_size=2048, intermediate_size=5504, hidden_layers=24, attention_heads=16), Transformer-2.7B (hidden_size=2560, intermediate_size=6880, hidden_layers=32, attention_heads=32).

Implementation Details. Unless otherwise specified, the patch size $K$ is 4. The context length for token-level training $2048$. For patch-level training, the context length is the patch size $K*2048$. The global batch size is $2M$ tokens, and the total number of training steps is $N=180000$. For patch-level training, the number of training steps is $N\lambda$, and then the model proceeds with token-level training for $N(1-\lambda)$ steps. After patch-level training, only the obtained model parameters are used for initialization, and all other states like the optimizer and learning rate scheduler are reset. We use the tokenizer of LLaMA2, whose vocabulary size is $32000$. Our models are optimized by the AdamW optimizer (Loshchilov & Hutter, ) with $\beta_{1}=0.9,\beta_{2}=0.95,\epsilon=1e-8$. The learning rate is $3e-4$ and the cosine learning rate schedule is applied with warmup of $2000$ steps. We use a weight decay of $0.1$ and gradient clipping of $1.0$, and no dropout is applied during training.

### 3.2 Main Results

**Table 1: Performance comparison of Transformers trained on the Pile dataset. $\lambda$ denotes the proportion of training data used for patch-level training, with the patch size $K$ fixed at 4. ‘PPL’ represents the perplexity score on the WikiText-103 test set. For zero-shot evaluations, we report the normalized accuracy across 6 NLP benchmarks. ‘Average’ means the average zero-shot accuracy.**
| Model Type | Cost | PPL | MMLU | HellaSwag | PIQA | WinoG | ARC-E | ARC-C | Average |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Transformer-370M | 1.0$\times$ | 10.9 | 22.9 | 40.8 | 67.5 | 53.1 | 44.3 | 24.7 | 42.2 |
| + Patch (${\lambda=1/2}$) | 0.625$\times$ | 10.6 | 23.5 | 42.0 | 67.9 | 52.1 | 46.1 | 25.6 | 42.9 |
| + Patch (${\lambda=2/3}$) | 0.5$\times$ | 10.7 | 23.7 | 41.1 | 68.0 | 51.9 | 46.0 | 24.2 | 42.5 |
| + Patch (${\lambda=4/5}$) | 0.4$\times$ | 11.0 | 23.3 | 40.5 | 67.5 | 51.7 | 44.9 | 24.5 | 42.1 |
| Transformer-780M | 1.0$\times$ | 9.2 | 24.4 | 48.5 | 69.0 | 55.4 | 49.0 | 26.7 | 45.5 |
| + Patch (${\lambda=2/3}$) | 0.5$\times$ | 9.1 | 24.1 | 49.1 | 70.6 | 54.8 | 51.3 | 28.2 | 46.3 |
| Transformer-1.3B | 1.0$\times$ | 8.2 | 23.9 | 54.5 | 71.2 | 57.3 | 55.1 | 28.9 | 48.5 |
| + Patch (${\lambda=2/3}$) | 0.5$\times$ | 8.2 | 24.3 | 54.1 | 71.6 | 57.8 | 55.6 | 30.4 | 49.0 |
| Transformer-2.7B | 1.0$\times$ | 7.1 | 25.3 | 62.2 | 74.3 | 61.5 | 61.2 | 34.3 | 53.1 |
| + Patch (${\lambda=2/3}$) | 0.5$\times$ | 7.2 | 25.4 | 61.9 | 74.9 | 62.4 | 61.9 | 34.6 | 53.5 |

Figure: Figure 4: Instruction-following abilities evaluated on MT-bench, a multi-turn question set.
Refer to caption: extracted/5851977/mtbench.png

We train a series of LLMs of varying sizes (370M-2.7B parameters) on the Pile dataset. To halve the training costs, we employ patch-level training with the settings of $K=4,\lambda=2/3$, and compare its performance with the conventional token-level training. For the Transformer-370M, we also explore other choices of $\lambda$ to evaluate its impact. Table 1 presents the performance comparison of the resulting models. Remarkably, our approach consumes only half of the compute and incurs almost no performance loss. It matches the baseline model in terms of perplexity and even demonstrates a consistent gain in zero-shot evaluations, raising the average accuracy by approximately 0.5%. The model performance is also influenced by the choice of $\lambda$. Within the range of values we set, a smaller $\lambda$ leads to better model performance but also entails larger training costs. A more detailed study on the hyperparameter $\lambda$ will be presented in Section 3.6.

We further conduct instruction fine-tuning using the Alpaca dataset by GPT4 to examine the impact of patch-level training on the model’s instruction-following ability. We evaluate our models using MT-Bench, a multi-turn question set, and present the experimental results in Figure 4. As can be seen, our approach maintains a similar instruction-following ability to the original models, with some experiencing a score decrease (Transformer-370M, Transformer-1.3B) and others showing an improvement (Transformer-780M, Transformer-2.7B), which can be viewed as regular variations.

Our primary motivation for patch-level training is to enhance the model’s knowledge acquisition efficiency. Interestingly, experimental results show that this approach can sometimes lead to performance improvements, which is beyond our initial expectation. We initially thought that the extended context length in patch-level training contributes to the improvements. However, when we decreased the context length during patch-level training from $KT=8192$ to $T=2048$ for Transformer-370M ($\lambda=1/2$), the model performance only experiences a slight decline (PPL $\uparrow$ 0.06, zero-shot accuracy $\downarrow$ 0.2), yet still surpasses the baseline, implying that context length is not the primary factor. We hypothesize that two other factors might be responsible for these improvements: firstly, the patch-level initialization could potentially serve as a form of regularization; secondly, by compressing tokens into patches, the model might more effectively recognize and capture long-range dependencies due to the reduced token distance.

**Table 2: Performance comparison of Transformers trained on 60B tokens for 6 epochs.**
| Model Type | Cost | PPL | MMLU | HellaSwag | PIQA | WinoG | ARC-E | ARC-C | Average |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Transformer-370M | 1.0$\times$ | 11.0 | 23.6 | 40.8 | 66.5 | 50.8 | 44.8 | 25.2 | 42.0 |
| + Patch (${\lambda=1/2}$) | 0.625$\times$ | 10.4 | 23.9 | 43.3 | 67.5 | 55.6 | 44.4 | 26.1 | 43.5 |
| + Patch (${\lambda=2/3}$) | 0.5$\times$ | 10.5 | 24.7 | 42.4 | 67.9 | 51.9 | 45.3 | 24.7 | 42.8 |
| + Patch (${\lambda=4/5}$) | 0.4$\times$ | 10.7 | 23.0 | 41.5 | 67.0 | 52.0 | 45.1 | 25.4 | 42.3 |

### 3.3 Multi-Epoch Training

Given that patch-level training consumes training data more rapidly, it is more data-hungry compared to token-level training. Consequently, it is essential to consider scenarios where training data is relatively limited and assess the performance of patch-level training when training data is reused for multi-epoch training (Muennighoff et al., ). We randomly extract a subset of 60B tokens from the Pile dataset and increase the number of training epochs to $N=6$. In this way, the model is first trained on patch-level for $N\lambda$ epochs, followed by $N(1-\lambda)$ epochs of token-level training.

The results presented in Table 2 demonstrate that patch-level training maintains its superiority in terms of training efficiency and performance on multi-epoch training. Intriguingly, when both consuming 360B tokens, patch-level training for multiple epochs even outperforms its single-epoch variant in Table 1. This unexpected advantage may stem from the synergistic effect of integrating patch-level and token-level training on the same data, which likely enhances model regularization. It also suggests that our approach can be data-efficient by initializing the model with patch-level training for one or multiple epochs, offering a promising direction for boosting model performance.

### 3.4 Scaling Properties

In the above, we have validated the effectiveness of patch-level training across several model sizes (370M-2.7B), using a training set of 360B tokens. However, state-of-the-art LLMs are generally trained on model sizes and datasets that are at least an order of magnitude larger than our settings. Therefore, it is crucial to know the scaling properties of patch-level training, i.e., how it performs when applied to larger training datasets and models.

In Table 1, we notice a trend of perplexity related to the model size: the performance advantage of patch-level training appears to decrease as the model size increases. Table 3 describes this trend more precisely, indicating that the model with patch-level training experiences smaller performance gains from the increase in model size. On the other hand, Table 4 presents the perplexity changes when maintaining a constant model size and varying the size of the training data. As the data size increases, the performance of patch-level training improves at a faster rate compared to the baseline model.

**Table 3: Perplexity scores when scaling the model size from 370M to 2.7B and training on the Pile dataset (360B tokens). ‘$\downarrow$’ indicates the perplexity reduction compared to the previous model size.**
| Model Size | 370M | 780M | 1.3B | 2.7B |
| --- | --- | --- | --- | --- |
| Transformer | 10.94 | 9.18 ($\downarrow$1.76) | 8.18 ($\downarrow$1.0) | 7.11 ($\downarrow$1.07) |
| + Patch (${\lambda=2/3}$) | 10.68 | 9.10 ($\downarrow$1.58) | 8.23 ($\downarrow$0.87) | 7.24 ($\downarrow$0.99) |

**Table 4: Perplexity scores of Transformer-370M when scaling the size of training data from 45B to 360B. ‘$\downarrow$’ indicates the perplexity reduction compared to the previous data size. The batch size is adjusted to maintain a consistent number of training steps.**
| Data Size | 45B | 90B | 180B | 360B |
| --- | --- | --- | --- | --- |
| Transformer | 12.50 | 11.71 ($\downarrow$0.79) | 11.28 ($\downarrow$0.43) | 10.94 ($\downarrow$0.34) |
| + Patch (${\lambda=2/3}$) | 12.84 | 11.80 ($\downarrow$1.04) | 11.17 ($\downarrow$0.63) | 10.68 ($\downarrow$0.49) |

This phenomenon can be explained from the perspective of knowledge transfer. As the data size increases, more training data is employed to adjust the model from patch-level to token-level, facilitating a smoother knowledge transfer process. However, an increase in model size implies a greater number of model parameters to be transferred to the token-level, which raises the level of transfer difficulty and necessitates more training data. Based on this explanation, patch-level training is better suited for scenarios with abundant training data.

Note that the above conclusions are drawn under the settings of $K=4,\lambda=2/3$, which may vary with changes in the patch size $K$ and the patch-level data fraction $\lambda$. At present, we have not identified a general scaling law for patch-level training that incorporates $K$ and $\lambda$. Instead, we have made some observations regarding their effects on model performance, which will be discussed in the following.

### 3.5 Effect of Patch Size K 𝐾 K italic_K

We investigate the effect of patch size under the settings of 90B training data, 370M model parameters, a batch size of 512K, and $\lambda=1/2$. The results are shown in Figure 5. Across different patch sizes, the loss curves for patch sizes $K=2$ and $K=4$ are nearly indistinguishable, while further increasing the patch size to 8 or 16 results in a certain performance decline. Despite this, these models still exhibit significant performance improvements when compared to the model trained from scratch, which does not benefit from the initialization of patch-level training.

Figure: Figure 5: Test losses of Transformer-370M w.r.t the number of processed tokens. Models are initialized by patch-level training with patch size $K$.
Refer to caption: x1.png

Overall, the patch size of $K=4$ strikes a favorable trade-off between training efficiency and performance. Considering that larger patch sizes can process more data at the same cost, we also experiment with patch-level training using $K=8$ on 90B tokens, which costs similar compute as $K=4$ on 45B tokens. Following this, both models proceed with token-level training on 45B tokens, and coincidentally, their loss curves are nearly identical. In this context, the advantage of $K=4$ lies in its data efficiency, as it achieves similar performance while consuming less data.

### 3.6 Effect of λ 𝜆 \lambda italic_λ

The hyperparameter $\lambda$ allocates the ratio of training data between patch-level and token-level training. A larger $\lambda$ results in more tokens being compressed into patches, leading to a higher acceleration rate, but it may also leave insufficient data to adjust the model to the token-level. In this section, we investigate the effect of $\lambda$ under the settings of 370M model parameters, a batch size of 512K, and a patch size of $K=4$. We consider two scenarios:

- 1.
Unlimited computational budget: We assess the impact of varying $\lambda$ while keeping the data size constant (90B tokens). The results are shown in Figure 7.
- 2.
Unlimited training data: We identify the optimal $\lambda$ under a fixed amount of computational budget (tokens + patches = 56.25B). For example, when $\lambda=1/2$, the size of training data should be 90B tokens, with 45B tokens being compressed into 11.25B patches. The results are shown in Figure 7.

Figure: Figure 6: Effect of varying $\lambda$ while keeping the data size constant.
Refer to caption: x2.png

Figure 7 shows that the model performance initially rises and later falls as $\lambda$ increases, with a turning point near $\lambda=1/4$. The performance improvements when $\lambda<1/4$ can be attributed to the inherent benefits of patch-level training, as analyzed in Section 3.2. When $\lambda$ exceeds $3/4$, further increasing $\lambda$ leaves insufficient data to adjust the model to the token-level, leading to a rapid decline in performance. Figure 7, on the other hand, shows that when computational budget are limited, the optimal value for $\lambda$ is around $2/3$. Note that these conclusions are specific to the current settings and should be used as a reference only. The optimal $\lambda$ may vary depending on factors such as data size and patch size. To determine the optimal value of $\lambda$ in any scenario, it is essential to establish the scaling law for patch-level training.

### 3.7 Effect of Architecture

In the current setup, the architecture of the patch-level model is identical to that of the token-level model, facilitating smoother model transfer but also leading to some concerns. For instance, the token-to-patch transformation and the prediction of the next patch do not take into account the order of tokens within a patch. It has been suggested that patch-level models might benefit from adopting architectures that are better suited for modeling and predicting patches, where the additional parameters introduced in such architectures could simply be discarded when transferring the model to token-level training.

We evaluate this strategy under the settings of 90B training tokens, 370M model parameters, a batch size of 512K, and $\lambda=1/2$. Specifically, we incorporate a linear mapping layer at both the input and output sides of the model. On the input side, the conversion of token embeddings into patch embeddings is facilitated through a linear mapping $w_{in}\in\mathbb{R}^{Kd\times d}$. On the output side, a linear mapping $w_{out}\in\mathbb{R}^{d\times Kd}$ is employed to transform patch-level representations back into token-level representations, followed by a softmax layer to obtain the probability distribution of each token. The effects of these two modules are detailed in Table 5.

**Table 5: Impact of architecture modifications in patch-level models. ‘+ ​InputMap’ and ‘+ ​OutputMap’ denote the incorporation of linear mappings at model’s input and output, respectively. ‘Patch PPL’ and ‘Token PPL’ are perplexities of the patch-level model and the final token-level model, respectively.**
| Model | Transformer | + ​InputMap | + ​OutputMap | + ​Both |
| --- | --- | --- | --- | --- |
| Patch PPL | 159.17 | 146.93 | 99.48 | 86.50 |
| Token PPL | 11.46 | 11.63 | 11.50 | 12.33 |

Overall, while these modifications are effective in reducing patch-level loss, they do not translate into benefits for the subsequent token-level training. Particularly, when both modifications are incorporated, although the patch-level loss is optimized to its lowest, the performance of the subsequent token-level model significantly declines. This suggests that there is no direct correlation between patch-level loss and the final performance of the token-level model, so a large loss during patch-level training does not imply ineffective learning. Therefore, we opt to preserve the original Transformer architecture and employ the most straightforward method for patch-level training.

### 3.8 Neuron Activation

In this section, we quantitatively explain why patch-level training leads to better learning efficiency from the perspective of neuron activation. The training of LLMs is essentially a process of embedding knowledge from the training set into the model’s parameters. During this process, the model employs all of its parameters to encode every token and updates the relevant parameters based on the gradient feedback. We argue that this is an inefficient process for large models, as the knowledge encapsulated in each token is only associated with a small subset of model parameters, resulting in a limited number of effectively activated and updated parameters.

We substantiate this by measuring the percentage of activated neurons for models of different patch sizes, as depicted in Figure 8. In the token-level model ($K\!=\!1$), only a small proportion of neurons are activated, suggesting that the model has enough capacity to handle text units with higher information density, thus indicating significant room for improvement in training efficiency. By grouping multiple tokens into a patch, the information density processed at each step is increased, which is manifested as increased neuron activation rates. From this observation, it becomes evident that patch-level training makes more comprehensive utilization of the model’s capabilities, therefore demonstrating higher training efficiency.

Figure: Figure 8: Percentage of activated neurons for models of different patch sizes. Output neurons of each model layer (FFN output) with an absolute value greater than 0.5 are classified as activated.
Refer to caption: x4.png

## 4 Related Work

Model Growth. Our approach draws inspiration from transfer learning, reducing training costs by transferring knowledge acquired at a lower training cost (patch-level) to a model with a higher training cost (token-level). A similar strategy has been employed in studies of model growth, which train large models at a relatively lower cost by progressively increasing the model size during training. For example, Gong et al. (); Yang et al. () improve the training efficiency by transferring knowledge from a shallow model to a deep model, where model layers are progressively stacked during training. Gu et al. () further proposes progressive compound growth, where the model grows at multiple dimensions during training, including the context length, model width, and the number of layers. Subsequent studies primarily focus on the initialization problem during the model growth process, i.e., how to expand the small model into a large one. Chen et al. (); Yao et al. () aim to achieve function-preserving growth (Chen et al., ) that the post-growth model have the same function as the pre-growth model, which intuitively ensures smooth knowledge transfer. Wang et al. (); Pan et al. () introduce learnable linear operators that linearly map the parameters of the small model to initialize the large model. Compared to model growth, patch-level training is more flexible and generalizable as it does not necessitate specialized model architectures or carefully crafted model mapping strategies. More importantly, patch-level training merely alters the information density of text units while preserving the full set of model parameters. Therefore, it enables the model to develop its complete capabilities without compromising its performance.

Multi-Token Prediction. Our approach improves training efficiency by concurrently predicting all tokens in the next patch. Similar attempts of multi-token prediction have been made in the past to improve the inference efficiency, including non-autoregressive generation (Gu et al., ) and speculative decoding (Stern et al., ; Leviathan et al., ; Chen et al., ). Non-autoregressive generation reduces the number of decoding iterations by generating all tokens at once, involving techniques such as knowledge distillation (Kim & Rush, ; Shao et al., ), training objectives (Shao et al., ; Ghazvininejad et al., ; Du et al., ; Shao & Feng, ; Ma et al., ), latent modeling (Kaiser et al., ; Ma et al., ; Shu et al., ), iterative decoding (Lee et al., ; Ghazvininejad et al., ; Gu et al., ), and expressive model architectures (Libovický & Helcl, ; Huang et al., ; Gui et al., ). Despite the reduction in decoding iterations, the computational demands (FLOPs) do not decrease and may even rise for certain structures, leading to increased training costs. Speculative decoding is a novel decoding paradigm for accelerating LLM inference. During each step of decoding, it drafts several future tokens efficiently and then verifies them in parallel. The model for generating draft tokens can either be a relatively small model (Leviathan et al., ; Chen et al., ) or a non-autoregressive model that generates multiple tokens in parallel (Cai et al., ; Lin et al., ; Fu et al., ; Li et al., ). Recently, pointed out that besides accelerating the inference, multi-token prediction also serves as an auxiliary training task to enhance the training signal. The key difference between our approach and these works lies in our utilization of multi-token prediction for reducing sequence length during training, with the goal of improving training efficiency. Regarding the model structure, we avoid introducing extra parameters and employ a single head for multi-token prediction.

Patch-Level Models. The concept of handling input data at the patch-level has emerged as a pivotal strategy for enhancing computational efficiency and capturing local features. Convolutional Neural Networks (CNNs, Lecun et al., ) are perhaps the earliest attempt of patch-level processing, utilizing kernel filters to extract local features from images. More recently, Vision Transformers (Dosovitskiy et al., ) have revolutionized image processing by employing CNNs to encode an image into non-overlapping image patches, thereby enabling Transformers to efficiently capture both local and global features. Similarly, speech models also rely on CNNs to compress high-frequency waveforms into hidden state representations (Baevski et al., ; Hsu et al., ), which can be interpreted as speech patches. Dosovitskiy et al. () suggested that the Vision Transformer should be fine-tuned at higher resolution, a strategy similar to ours, which improves data precision during the later phases of training. In subsequent studies, the impact of patch sizes has also garnered attention. Beyer et al. () proposed to randomize the patch size at training time, which creates models capable of handling diverse patch sizes and adapting to different compute budgets. Further exploration by Anagnostidis et al. () revealed that training with larger patch sizes is generally more computationally efficient. This insight led to a training strategy that progressively decreases the patch size, enabling more efficient training of vision Transformers. For textual data, characters, the basic building blocks of text, can be downsampled into more compact representations (Clark et al., ) or merged into tokens (Sennrich et al., ; Kudo & Richardson, ). Recently, there have been attempts to further compress tokens into patches, with the model directly processing the patch sequence (Nawrot et al., ; Mujika, ; Yu et al., ; Ho et al., ). However, it remains necessary to upsample the patch representation and input it into a token-level autoregressive model for the likelihood inference.

## 5 Conclusion

This paper introduces patch-level training, an efficient training approach for large language models, 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, the model reads training data in patches and learns to predict the next patch. Following this, a small amount of training data is utilized to adjust the model to the token-level. Experimental results show that this approach can cut LLM training costs by 50% while maintaining comparable performance.

Yet, our exploration of patch-level training is still in its infancy, and advancements in the following directions could further enhance this methodology: assessing the scalability of patch-level training by evaluating its performance on larger models and datasets; establishing an empirical scaling law for patch-level training, ideally incorporating both $K$ and $\lambda$; developing advanced training techniques to accommodate larger $K$ and $\lambda$, thereby pushing acceleration rates to a higher level; further investigating the potential of patch-level training in multi-epoch training; exploring the applicability of patch-level training to other data modalities, such as images, speech, and video.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: N/A
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
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.
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
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?
soundness: 3
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.)
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).
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.

---

confidence: 5
contribution: 4
initial_rating: 8
presentation: 3
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.
soundness: 4
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
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.
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.

Meta-review: 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.

Decision: Accept (Spotlight)