We thank you for the comments, and acknowledging the innovation of trading inference efficiency and parameters, the efficiency of Block Transformer, solid experiments, and good paper organization. We address weaknesses below.

.

W1. Bigger model size

Despite bigger model size, Block Transformers achieve higher throughput and lower total memory usage during inference. For training, it is possible to uptrain existing vanilla models with ~10% of training steps, and training costs can be cheaper under constrained inference budget.

Inference case

• Total memory usage: Block Transformers can achieve significantly higher batch sizes under same hardware, as the KV cache memory usage per sample of Block Transformers is 3—6x times lower, despite the larger size (Table 2). Note that the KV cache of a single sample can be as large as the parameters of the entire model (Table 2).
• Throughput: our main results show significant gains in throughput over vanilla models under same hardware constraints.
• These advantages extend into the multi-GPU inference scenario. Tensor parallelism can further reduce per-GPU model parameter memory and leaves more room for more KV cache, which is advantageous for the Block Transformer.

Training case

• IsoFLOP analysis in Section 3.6 shows that Block Transformers can achieve better performance and significantly higher throughput using the same training FLOPs as vanilla models when constrained by inference budget (we discuss further in W4).
• Section 3.7 shows that it is possible to uptrain existing vanilla models into Block Transformers, significantly reducing development costs. Uptrained models approach the performance of model pre-trained from scratch, using just 10% of training steps (Figure 5a and Appendix M).

.

W2. Doubts on inference efficiency comparison

We point out that Block Transformer does not “only save attention computation”, but have a wide range of advantages stemming from our architecture design. Please refer to Section 2 and Appendix D. We also provide a visual recap of advantages in Figure 1 of the attached PDF.

We re-iterate four key advantage of our architecture:

1. Attention + FFN Token decoder does not need to prefill prompt tokens.
2. Attention KV cache IO at the token decoder is reduced by $L/L_B=2048/4=256\times$.
3. Attention + FFN Block decoder operates at the block level, reducing overall costs by $L_B=4$. KV cache IO cost is reduced quadratically.
4. Attention + FFN Overall KV cache storage is reduced by $L_B=4\times$ in the block decoder and nearly eliminated in the token decoder. This enables higher batch sizes and thus higher compute utilization.

We also empirically pinpoint these benefits in new detailed measurements in Table 1 of the attached PDF. These show that Block Transformer speeds up all segments of computation relative to a loss-comparable vanilla model: {attention, FFN} operations at {lower, upper} layers during {prefill, decoding} stages under {prefill, decode}-heavy settings.

.

W3. No evaluation of long-sequence modeling capability

Figure 4b of Section 3.5 demonstrates the long-context modeling capabilities of Block Transformers. Our models achieve lower loss by utilizing more context up to 2K tokens on the PG19 test dataset. This is a standard benchmark used to evaluate long-context modeling capability [1, 2].

We also extend this to 8K context in Figure 2 of attached PDF, with new models pretrained on 8K context for ~30B tokens. We show that the Block Transformer (1.2B) achieves lower loss at all token positions relative to a comparable vanilla baseline (300M), despite achieving significantly higher throughput (Appendix Figure 10).

.

W4. Scaling property

In Appendix A, we discuss that rigorous Chinchilla-style scaling study is infeasible as it requires repeated training runs with appropriate LR schedules for each FLOP budget [3].

However, expanding on W1, we do discuss in Section 3.6 why our isoFLOP analysis (which is in the overtraining regime for the vanilla model) is relevant under current trends—where small models are significantly overtrained to maximize performance at a given inference budget [4]. I.e., to achieve maximize performance using plentiful training FLOPs under a tight inference budget with vanilla models, we must resort to training a small model into a highly suboptimal training regime. Since Block Transformers have significantly smaller inference costs, we can use larger models which can achieve higher better performance given the same training FLOPs, under inference budget constraints. This is what is shown in Figure 4c.

.

Q2. FlashAttention and FlashDecoding

Existing experiments: all models in our paper are based on GPTNeoX using MHA on Huggingface and used eager attention implementation for inference, due to limited support for FlashAttention2 during the time of submission*.

FlashDecoding [5]: New measurements using FlashAttention2’s FlashDecoding kernel* show a similar pareto frontier as our existing results in Figure 3 of attached PDF. Please compare with Figure 2 of main paper. We find that FlashAttention2 achieves 8-41% throughput gains in vanilla models and +8% to +31% in block models, in the prefill-heavy setting. Gains are diminished in larger models. We observe +16% to -37% speedup across models in the decode-heavy setting, without a coherent trend.

.

References

[1] Zhang, Zhenyu, et al. "H2o: Heavy-hitter oracle for efficient generative inference of large language models."

[2] Xiao, Guangxuan, et al. "Efficient streaming language models with attention sinks.”

[3] Hoffmann, Jordan, et al. "Training compute-optimal large language models.”

[4] Touvron, Hugo, et al. "Llama: Open and efficient foundation language models.”

I appreciate authors' thorough and patient response. I have a better understanding of the experiment setting. However, I still have some of my previous concerns:

1. Measuring the perplexity on PG19 to evaluate the long-sequence modeling capability is still not a good experiment. Since some subsequent works show that H2O and Attention Sink are not good at some long context tasks, including Needle-in-a-Haystack, LongBench, and ZeroScrolls, I don't agree it's a standard benchmark.

I believe it is a novel and insightful work. But there is still space for improvement. In a nutshell, the paper will look much stronger if the experiment is on a "modern" setting, including the latest evaluation settings and kernel techniques. I will keep my initial score.

I'm pleased to see that Block Transformer can effectively retrieve global information contained within the 2K context length, that is a strong indicator for long-context capability. I will increase my score to 6.

We thank you for the comments, and we are encouraged that you pointed out the trustworthiness of our extensive experiments with modern setup, and easy-to-follow writing. We address weaknesses below.

W1. Difference between Block Transformer and Funnel Transformers

Major difference in key aspect—local attention: we respectfully disagree that our Block Transformer is similar to existing pooling transformers, including Funnel Transformer [1], Hourglass [2], and CANINE [3]. Our approach is fundamentally different, as it applies local attention in the token decoder instead of maintaining global attention with pooling and up-sampling.

Pooling transformers achieve speedup by pooling as the depth increases and then upsampling in the upper layers. However, they maintain global attention throughout all layers. In contrast, while global-to-local modeling applies a pooling operation at the token level, a primary feature is the locality of attention in the token decoder (upper layers). The token decoder applies attention only within the local window, leading to significantly higher compute utilization.

We will include this discussion in our related works section.

Refer to advantages of this local attention (and more) in our visual recap in Figure 1 of attached PDF , and details in Section 2.4 and Appendix D.2 . We also empirically pinpoint the benefits of the locality aspect (Table 1 of attached PDF), accounting for ~90% reduction in walltime at the upper layers

.

W2. Pareto-Frontier of decoding throughput is only improved with batch sizes greater than 32

We respectfully disagree, and would like to emphasize that trends in Figures 7, 8 show the throughput gains increases as the model size increases, even at smaller batch sizes

In fact, Block 1.2B already surpasses Vanilla 300M in both prefill-heavy and decode-heavy throughput at batch size 1. This is because KV cache saturates GPU memory even at small batch sizes, with larger models.

Below, we show the throughput of models above 1B parameter with batch size 1 and 2. The advantage of Block 1.2B against the loss-equivalent Vanilla 300M model widens as batch size increases. We expect that Block 6.4B will perform between Vanilla 1.2B and 2.5B, and find that Block 6.4B is faster than Vanilla 1.2B. The gap also widens as batch size increases.

.

W3. Block Transformer requires larger model size, and costs more training compute and memory

To mitigate model development costs (including training), it is possible to uptrain existing vanilla models with just ~10% of training steps. We also find that training costs can be cheaper under constrained inference budget.

• Uptraining existing vanilla models into Block Transformers approaches the performance of models pre-trained from scratch, using just 10% of training steps (Section 3.7, Figure 5a, Appendix M).
• IsoFLOP analysis under in Section 3.6 shows that Block Transformer can achieve better performance and significantly higher throughput using the same training FLOPs as vanilla models. Note that this assumes overtraining to fit a fixed budget constraint, in line with recent trends set by models such as Gemma and Llama [4] (Section 3.6).

.

Q1. Do baselines use FlashAttention2 and MQA/GQA?

Existing experiments: all models in our paper are based on GPTNeoX using MHA on Huggingface and used eager attention implementation for inference, due to lack of support for FlashDecoding during the time of submission.

FlashDecoding [5]: New measurements using FlashAttention2’s FlashDecoding kernel* show a similar pareto frontier as our existing results in Figure 3 of attached PDF. Please compare with Figure 2 of main paper. We find that FlashAttention2 achieves 8-41% throughput gains in vanilla models and +8% to +31% in block models, in the prefill-heavy setting. Gains are diminished in larger models. We observe +16% to -37% speedup across models in the decode-heavy setting, without a coherent trend.

GQA: We expect that both Vanilla and Block Transformers will benefit from GQA, in terms of the attention operations.

• Q. what if FFN becomes the deciding factor after applying GQA?
• A. new measurements show that Block Transformer is significantly faster than vanilla in both attention and FFN operations (Table 1 in attached PDF).

Q3. Maintaining our claim

We made this claim because we were the first to identify the locality in the token decoder as the source of significant inference-time benefits, in the context of global-to-local modeling (coarse global attention in lower layers and fine local attention in upper layers).

We will revise the sentence to more clearly emphasize our recognition of the benefits of the local module in global-to-local modeling.

.

References

[1] Dai, Zihang, et al. "Funnel-transformer: Filtering out sequential redundancy for efficient language processing."

[2] Nawrot, Piotr, et al. "Hierarchical transformers are more efficient language models."

[3] Clark, Jonathan, et al. “CANINE: Pre-training an Efficient Tokenization-Free Encoder for Language Representation.”

[4] Touvron, Hugo, et al. "Llama: Open and efficient foundation language models.”

[5] Dao, Tri, et al. "Flash-Decoding for long-context inference”

We thank you for the comments, and we are encouraged that you pointed out the improved efficiency of our Block Transformers with extensive experiments. The weakness of our paper is discussed as follows.

W1. Difference between Block Transformer and MEGABYTE.

We have discussed the main differences to previous works related to global-to-local modeling in Appendix C.1. We believe our contributions are novel to Megabyte. Here’s why:

1. Firstly, the primary goal of our proposed architecture, which utilizes global-to-local language modeling, is clearly different. While MEGABYTE (including most hierarchical transformer works) focuses on efficient pretraining, we mainly focuses on efficient inference. MEGABYTE aimed to reduce training time by minimizing FLOPs through global-to-local modeling. They optimized architectures under a fixed FLOPs budget, leading them to favor a model with a six times larger global module compared to the local model.

   In contrast, we studied global-to-local modeling with an emphasis on inference throughput in autoregressive LMs. Specifically, we analyzed throughput trends based on the block length and parameter allocation ratio (refer to Figure 15) , and our findings reveal that increasing the size of the token decoder (local model) is beneficial for improving inference speed. This interpretation is completely overlooked and under-explored in MEGABYTE, which argue that “Many of the efficiency advantages of the MEGABYTE design could be realized with the global model alone”, thus use a significant small local model. This results in a remarkable speedup of up to 20x, contrasting their reported 1.4x improvement. New detailed measurements in Table 1 of the attached PDF empirically pinpoints the benefits of the locality aspect, accounting for the ~90% reduction in walltime at the upper layers compared to vanilla models.

2. Second, as you described, Block Transformers use subword inputs, while MEGABYTE uses byte-level inputs. This enables us to employ an uptraining strategy with initialization techniques that fully leverage existing subword-level language models (refer to Section 3.7 and Figure 5a). Specifically, we demonstrate that with only 10-20% of uptraining, we can almost match the performance of the original model. This is a significant contribution to future research, as it enables the conversion of high-performing LLMs into inference-specialized Block Transformers with minimal additional training cost. Additionally, further exploration of initialization methods could lead to even greater performance improvements and a reduction in the optimal parameter size of Block Transformers.

We would also appreciate acknowledgement for our novel findings and insights regarding global-to-local language modeling, such as the diverse conclusions drawn from analyzing the relationship between block length and parameter allocation ratio in terms of perplexity and throughput (refer to Section 3.3 and Section 3.5), extensive ablation studies on various components, especially token decoder components (refer to Section 3.4), IsoFLOP analysis with the inference speed budgets compared to vanilla transformers (refer to Section 3.6), or exploring the information contained in context embeddings, which has never been addressed in prior research (refer to Appendix P and Table 5).

Consequently, we strongly believe that our findings and contributions will provide valuable insights for future research on inference-optimized global-to-local language modeling.

We thank you for the thoughtful feedback. We are encouraged that you found our research direction interesting and throughput improvement significant. The weaknesses of our paper are discussed as follows.

.

W1. More parameters are needed to achieve the similar perplexity.

KV cache IO and memory size typically impacts generation speed during the decoding process. In standard transformers, larger parameters necessitate larger KV cache sizes, negatively affecting both speed and memory consumption. Conversely, our Block Transformer mitigates this issue by grouping tokens into blocks and performing local modeling on each block embedding (see memory and IO comparison in Table 1).

Therefore, despite having two or three times large parameters, the significantly reduced KV cache sizes enable the Block Transformer to demonstrate up to 20 times faster generation speed, taking into account the total time spent on parameter loading or KV cache loading. Our model also handles much larger batch sizes, making it a more efficient solution for real-world applications. In Table 1 of the attached PDF, we have summarized the actual speed improvements achieved in the attention and FFN operations of both the block and token decoder. Additionally, Figure 1, in the attached PDF, illustrates which model elements contribute to increased throughput despite having more parameters.

It is also important to note that our Block Transformer is not yet a fully optimized architecture, meaning there is still room for improving performance while maintaining its speed advantage. For example, as confirmed by attention scores (refer to Appendix O), incorporating multiple, salient block embeddings in token decoder could substantially enhance performance, potentially reducing the required parameter sizes of Block Transformer. Meanwhile, due to the nature of local modeling, slightly increasing the context length in local modules would have minimal impact on the actual generation time.

Moreover, our token-level modeling structure allows us to uptrain from a well-pretrained checkpoint (refer to Figure 5a and Appendix M). Further exploration of initialization methods to effectively leverage pretrained models could potentially reduce the parameter size of the Block Transformer to achieve the same perplexity, while significantly reducing the training time as well.

.

W2. Scaling Models up to 7B or 13B parameters.

Performance scaling: we acknowledge the important of scaling studies. Unfortunately, due to the significant computational resources required for pretraining from scratch, we were unable to verify our findings at larger scale such as 7B or 13B. However, as a proof of concept, we have successfully demonstrated that our proposed modeling approach can achieve similar perplexity across six different scales up to ~1B parameters. Based on our extensive studies, we believe that scaling Block Transformers beyond 7B parameters will still yield compelling perplexity.

Inference throughput scaling: independent of perplexity, we compared the throughput of the Vanilla and Block Transformers scaled up to 7B parameters using random initialized weights. We present the maximum throughput (1k tokens/sec) of models not included in Table 2, in prefill- and decode-heavy scenarios. We find that our Block Transformer with 6.9B parameters is still faster than that of a 160M vanilla model. Assuming that Block Transformer 6.9B performs between that of Vanilla 1.4B and 2.8B, we can still expect at least a 6x speed-up (by comparing with vanilla 1.4B).

.

W3. Lack of experiments on longer contexts like 8K or 16K.

To further support the effectiveness of our proposed block language modeling in capturing full context, we conducted experiments with an 8K context length (please refer to Figure 1 in the attached PDF). However, due to limited computational resources, we pretrained only a 70M parameter vanilla model and a 170M parameter Block model. Following prior work [1,2,3] that used token position-wise perplexity on the PG19 dataset to demonstrate the utilization of global information in long contexts, we evaluated our Block Transformer in the same manner with 8K context length (refer to Figure 4b of the main paper for 2K context window). Even with a extended 8K context window, our models effectively utilized the full context, showing a decreasing loss trend as token position increased, similar to the vanilla model. Besides, consistent with Table 2 in the main paper, the 170M block model outperformed the 70M vanilla model in terms of perplexity.

It is also worth noting the robustness of our proposed block language modeling in leveraging full context, regardless of the block length. As shown in Figure 4b of the main paper, even when varying the block length from 1 to 8, the models exhibited the same loss slope with respect to token positions. This indicates that our block language modeling can capture the full context robustly, regardless of the degree of compression applied to context representations.

.

Reference

[1] Yu, Lili, et al. "Megabyte: Predicting million-byte sequences with multiscale transformers."

[2] Xiao, Guangxuan, et al. "Efficient streaming language models with attention sinks."

[3] Zhang, Zhenyu, et al. "H2o: Heavy-hitter oracle for efficient generative inference of large language models."