Block Transformer: Global-to-Local Language Modeling for Fast Inference

Summary

This paper proposes a Block Transformer architecture which adopts hierarchical global-to-local modeling to mitigate the computational cost and KV cache memory of Self Attention. Block Transformer isolate the global modeling with three blocks: Embedder, Block Decoder, and Token Decoder. Embedder encodes block information for auto-regressive modeling in Block Decoder. The aggregated information is plugged into Token Decoder for final token-level decoding. Block Transformer shows 10-20x inference improvement under the similar perplexity than traditional global Transformer. This paper gives a detailed ablation about the architecture design of Block Transformer. Besides, the global modeling capability and uptraining strategy is also discussed.

Strengths

1. The paper is well organized and the experiments are solid and comprehensive. 2. Block Transformer trades inference efficiency with model parameters, which is a bold innovation. 3. Block Transformer achieves low-cost inference compared with standard Transformer. 4. The architecture analysis and ablation studies show the effectiveness of Block Transformer design.

Weaknesses

1. Block Transformer uses much bigger model size to compensate for the performance loss. There are countless problem for that, including training partitioning and inference infrastructure. 2. The actual inference efficiency comparison is doubtful. For example, in Table 2, The Decode throughput of Block Transformer is much bigger than standard Transformer under same model size. Since Block Transformer only saves Attention computation and the FFN computation stays the same, I'm confused that why attention computation occupies almost most of the overall computation. 3. The long-sequence modeling capability is not evaluated. Since Block Transformer squeezes the context representation, it is questionable that if Block Transformer can retrieve the global context information. I think some common long-sequence experiments will help, e.g., Needle-in-a-Haystack. 4. The scaling property is not discussed good enough. For example, in Figure 2, there are some results in different model sizes. However, there is not a "scaling law" for Block Transformer. Besides, the scaling does not look promising with my human eye in Figure 2.

Questions

1. Concerns in the Weaknesses part. 2. I'm curious about the inference experiment setting. Are FlashAttention and FlashDecoding techniques used for vanilla Transformer or Block Transformer? I believe it is already a necessary part in 2024 year.

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

The limitation is discussed.

Summary

The authors introduced the Block Transformer architecture to address the self-attention bottleneck. This is achieved by grouping input tokens into fixed-size blocks and applying self-attention at a corser level throughout the model. At the output layer, a token decoder predicts individual tokens from the block. The authors found that this hierarchical global-to-local modeling approach results in 10 to 20 times faster inference compared to vanilla transformers with similar perplexity.

Strengths

- The topic of improving the efficiency of LLMs and making them more affordable is crucial and timely. Furthermore, as sequences scale, the self-attention bottleneck increases. - The authors conducted extensive experiments and ablations with a modern setup. The evaluation includes perplexities and zero-shot downstream tasks. The models were trained for a significant number of tokens (300B, which is more than an epoch on the Pile), making the result more trustworthy. - The paper is well-written and easy to follow.

Weaknesses

- The concept appears similar to that of the Funnel Transformer [1], with the main exception that the aggregation happens only once at the token level. - The Pareto front of the decoding throughput is only improved with batch sizes greater than 32, which may not often be the case. - The higher inference throughput comes at the cost of more training compute and memory. The proposed methods perform worse than the vanilla model at an equivalent size. [1] Dai, Zihang, et al. "Funnel-transformer: Filtering out sequential redundancy for efficient language processing." Advances in neural information processing systems 33 (2020): 4271-4282.

Questions

- Is the baseline using FlashAttention2 [2] and MQA/GQA [3] ? These two modifications have become standard and significantly reduce the bottleneck of the attention. - Could you explain the differences between your methods and the Funnel Transformer? - Given the Funnel Transformer and the many other sparse attention, do you maintain your claim that you "are the first to recognize the central role and inference-time benefits of both global and local modeling in autoregressive transformers, particularly the significance of local modules"? [2] Dao, Tri. "Flashattention-2: Faster attention with better parallelism and work partitioning." arXiv preprint arXiv:2307.08691 (2023). [3] Ainslie, Joshua, et al. "Gqa: Training generalized multi-query transformer models from multi-head checkpoints." arXiv preprint arXiv:2305.13245 (2023).

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes.

Summary

The paper introduces the Block Transformer architecture, which aims to improve inference speed in autoregressive language models by adopting a hierarchical global-to-local approach. The architecture separates the global context modeling into lower layers and local detailed interactions into upper layers, thus reducing the self-attention bottleneck. The authors demonstrate significant improvements in inference throughput without compromising perplexity.

Strengths

1. The experiments in this paper are extensive, covering a variety of model parameters. 2. The proposed Block Transformer demonstrates improvements in computational efficiency, which is crucial for scaling up to longer sequences.

Weaknesses

**Major Weakness** As I understand it (please correct me if I'm wrong), the primary difference between Block Transformer and MEGABYTE [1] is whether the input is a token or a byte. The architecture of Block Transformer is nearly identical to that of MEGABYTE, which significantly limits the novelty and contribution of this work. [1] MEGABYTE: Modeling Million-byte Sequences with Multiscale Transformers

Questions

None

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

1

Limitations

As discussed above.

Summary

This paper introduces Block Transformer, which is a new architecture that adopts hierarchical global-to-local modeling to autoregressive transformers to mitigate the inference bottlenecks brought by applying self-attention on the global context. In detail, Block Transformer mainly includes three different components: (1) Embedder, which aggregates each block into an input block embedding; (2) Block decoder, which applies self-attention on the full sequence of blocks (rather than tokens) to model global context; (3) Token decoder, which applies self-attention on the sequence of tokens within each block to model local context and decode individual tokens. Evaluation shows that the Block Transformer architecture demonstrates significant gains in inference throughput compared to vanilla transformers with similar perplexity.

Strengths

- The paper explores an important and interesting research direction. - The improvement on inference throughput achieved by Block Transformer is significant. - The paper is generally well-written.

Weaknesses

- Block Transformer needs two or three times more parameters than vanilla transformers to achieve similar perplexity. - It is unclear that, after scaling up vanilla transformers to 7B or 13B level, whether Block Transformer can still achieve similar perplexity with two or three times more parameters. - More evaluation is required to demonstrate that Block Transformer can effectively leverage full context. While the paper evaluates the perplexity of token positions within a 2K context window to show that Block Transformer can effectively leverage at least 2K tokens of context, experiments on longer contexts that are no shorter than 8K or 16K is also important to show that Block Transformer can indeed effectively leverage global information.

Questions

None beyond the above.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The limitation section of this paper is rather comprehensive.