Author response
> Unclear conclusions about the effect of architecture. Highlight some use cases where subsequent token level training does not work very well.
We thank the reviewer for this valuable suggestion. In the architecture ablation, we observed that it is crucial for at least one end (input or output) of the patch-level model to remain consistent with the token-level model, which acts as an anchor to align their representations. Overall, the best practice is to keep the model architecture unchanged, as it facilitates smooth knowledge transfer between the two training phases. We have updated the manuscript to emphasize these concepts in Section 2 and Section 3.7.
> No numbers showing the actual speed improvements.
We thank the reviewer for raising this concern. First, we would like to address a misunderstanding about the computational process of patch prediction, which actually only needs one softmax operation to predict the next patch. Below is the key modifications in our code. On the input side, we average the token embeddings to obtain the patch embedding, with a negligible complexity of $\mathcal{O}(KTd)$. On the output side, we first perform a softmax operation to predict the log-probability of next patch, and then calculate cross-entropy losses for each token in the next patch. The added computational burden for this loss computation remains modest at $\mathcal{O}(KT)$.
```python
# Model input
num_patches = seq_length // self.patch_size
inputs_embeds = inputs_embeds.view(batch_size, num_patches, self.patch_size, -1).mean(2)
position_ids = position_ids[:, :num_patches]
...
# Model output
shift_logits = logits[..., :-1, :].reshape(-1, self.config.vocab_size)
shift_labels = labels[..., self.patch_size:].reshape(-1, self.patch_size)
loss = 0
log_probs = F.log_softmax(shift_logits, dim=1)
for i in range(self.patch_size):
loss = loss + F.nll_loss(log_probs, shift_labels[:, i])
loss = loss / self.patch_size
```
The reviewer is correct that the practical speed-up achieved through patch-level training falls short of the theoretical $K\times$ improvement, primarily due to overheads of data loading and gradient synchronization. Compared with the basic setting (patch_size=1, block_size=2048, per_device_train_batch_size=4, gradient_accumulation_steps=8, num_gpu=8), increasing the patch size to 4 and the block length to 8192 results in a speed-up of approximately 3.8$\times$, with some efficiency loss due to data loading. In practice, the gradient accumulation steps should be reduced to 2 to keep the total batch size constant, which brings the speed-up down to around 3.5$\times$ due to more frequent gradient synchronizations. Consequently, the actual runtime is approximately $\frac{2}{3} \times \frac{1}{3.5} + \frac{1}{3} \approx 0.523\times$, which slightly exceeds the theoretical cost of $0.5\times$. We have revised our manuscript to include a detailed analysis of the practical speed improvements in Appendix B.