Block-State Transformers

Summary

The authors present a new long-range transformer architecture by incorporating SSMs. This novel model outperforms several established baselines, such as Transformer XL, Block Recurrent Transformer, and Sliding Window Transformer, in terms of cost-effectiveness trade-off for tasks involving long-document or code modeling.

Strengths

1) Well-motivated. Long-range modeling is becoming increasingly important for LLM community. 2) Good results on language modeling (PG19, arXiv, Github) 3) The writing is clear and effectively conveys the ideas and findings. 4) The model design is intuitive and well-reasoned. The inclusion of both local full attention and linear components to handle long sequences is a sensible approach.

Weaknesses

1) The scale is too small. For language modeling, based on the success of LLM, we always expect good scalability. This paper only conduct experiments up to 380M params, which are far from many emergent abilities threshold. When scaling up, many inductive bias would become useless[1]. 2) More insightful experiments beyond language modeling are required. For instance, as the authors mentioned in the limitation section, what about the results on Long Range Arena? 3) Any case study about how this model captures long-range dependency? Why this model is indeed better? It seems that putting one efficient attention layer with linear or sub-linear complexity before self-attention should work similarly. 4) Some strong baselines like CoLT5[2] are missing. [1] Scaling Laws vs Model Architectures: How does Inductive Bias Influence Scaling? [2] CoLT5: Faster Long-Range Transformers with Conditional Computation I enjoyed reading this idea but I believe the missed experiments above, especially (1), would be highly desirable. Without a set of experiments about scaling the model up, I cannot agree this paper is useful enough.

Questions

1) Did the authors observe any instable training process? And again, would it become instable when scaling up?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

See weakness

Summary

This paper focuses on combining two efficient techniques for long-range modeling: state-space models (global contextualization) and block-recurrent transformers (local contextualization). In particular, they propose two different approaches, the first uses SSMs to output contexts for multiple heads (multi-head), and the second concatenates the last entries from the previous window to form a combined context state (multi-filter). Evaluation is performed on three language modeling datasets that outperform block-recurrent transformers in perplexity and is much faster when compared layer-wise.

Strengths

S1. Exploring different ways to combine SSMs and block-recurrent models to improve efficiency is a compelling direction. SSMs offer a parallelizable way to capture long-term information and avoid sequential computation in block-recurrent models. The results of this study should be of interest to researchers that study architectures that capture local and global information. S2. The evaluation even though it focuses mainly on comparison with block-recurrent models and SSMs on three language modeling tasks, it is thoroughly described and well-executed.

Weaknesses

W1. Even though the paper is mainly empirically driven, the delivery lacks a comprehensive and diverse set of evaluations to demonstrate the effectiveness and limitations of the method. W2. Experiments are targetting language modeling on three tasks but there is no experiment that measures the long-range capabilities of the model. There are several long-context classification benchmarks that the authors can use in addition to language modeling: LRA [1], MuLD [2], and CAB [3]. W3. The method design makes specific assumptions about the hardware to be employed and bases its evaluation on it; e.g. efficiency comparisons are made per layer. It's not explored to what extent the benefit remains when comparing training time/speed vs quality for the whole model and evaluation is performed on typical accelerators. That reduces the practical impact in my view. W4. A study regarding the behavior of the model when increasing the model size is missing. Scaling aspects are important to consider when making claims about outperforming transformers. [1] https://arxiv.org/pdf/2011.04006.pdf [2] https://arxiv.org/pdf/2202.07362.pdf [3] https://arxiv.org/pdf/2210.07661.pdf

Questions

Q1: Evaluating methods that model long context typically involves more tasks than language modeling. What is the level of confidence that proposed models perform to other tasks such as long context classification and text generation? Q2: It would be interesting to measure the time it takes for every method to converge even within a fixed time budget. Have all methods converged in the fixed training budget in Table 1 and do they have any differences worth discussing?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

Discussion about the limitations of the proposed method would be useful, I'd suggest talking about scaling behavior, performance on conventional accelerators, and generalizability to long-context tasks.

Summary

State space models (SSMs) perform well on modeling long-range dependencies with good efficiency scaling, but on language modeling, transformers still outperforms SSMs. This paper tries to combine the best of both worlds and proposes a hybrid model, Block-State Transformer, which combine SSMs’ capacity on long range modeling and Transformer’s ability on modeling local context. The input sequence is split to multiple smaller segments. For each segment, transformer layer will do a self attention on this token embeddings and cross attention to the output of SSMs. Their experiments show that on language modeling, their approach achieve reasonable speedup with comparable performance to Transformers.

Strengths

1. The proposed combination of SSMs and transformers allow the model to exploit advantages of two powerful methods while avoiding their drawbacks. 2. The SSMs used in the proposed method can be swapped to different SSMs making it possible to enjoy the advancement on SSMs field. 3. The proposed method give similar performance on language modeling compared to Transformers.

Weaknesses

1. There was already existing work on SSMs that achieve similar performance on language modeling (https://arxiv.org/abs/2212.14052) compared to Transformers. The authors should include a discussion and comparison to the relevant work. 2. The evaluations are performed on language modeling for 4096 length sequences. On this setting, there are a lot of strong transformer baselines with efficient self-attention designs. It would be good if the authors can provide an empirical comparison with these baselines. 3. Perplexity is only one indicator of how the language models performance. To get more precise understanding of performance, it would be good to include a comparison on downstream tasks. 4. Code is not available.

Questions

The main concerns are listed above. There are two more questions: 1. On line 235, the 6,966,499 English language words seems to be a typo according to dataset statistics on https://github.com/deepmind/pg19. 2. What are the latency of one forward step and what about memory consumption?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes.

Summary

This paper proposes block-state transformers, a method to combine state space models with transformers for language modeling. The paper evaluates block-state transformers on PG19 and arxiv math and finds promising results.

Strengths

Combining state space models and Transformers is an interesting idea worth exploring. The presentation of the paper is clear. The explanation of state space models, which can be quite complex, is very clear. The evaluation hits the right notes in terms of the major questions to ask.

Weaknesses

The evaluation is missing many recent works combining state space models and attention in various ways. The claim that SSMs do not match Transformers on language has not been true for a while. Most of these methods were released significantly before the NeurIPS deadline and are critical to compare against for evaluation. * Mega [1] combines attention and state space models. * BiGS [2] is a new SSM-based architecture that matches Transformers in language. * H3 [3] combines SSMs and attention in alternate layers. * Hyena [4] removes attention completely and replaces it with a convolution-based layer (similar to an SSM). Confusingly, many of these works are cited in the paper - and ideas from the papers are used extensively in the methods proposed (e.g., "BST:{SH,MF}" uses the structure from H3 and Hyena without comparing against those architectures as baselines). Using the ideas from these papers without comparing against them makes it difficult to understand how this method compares against previous methods and where the performance improvement comes from. Performance is also hard to evaluate compared to standard models such as Transformers (GPT-Neo [5], Pythia suite [6]). TransformerXL is an older model that is not trained as well as modern Transformer-based LLMs. [1] https://arxiv.org/abs/2209.10655 [2] https://arxiv.org/abs/2212.10544 [3] https://arxiv.org/abs/2212.14052 [4] https://arxiv.org/abs/2302.10866 [5] https://github.com/EleutherAI/gpt-neo [6] https://github.com/EleutherAI/pythia

Questions

How does BST compare to the architectures listed in the weaknesses section? It is important to compare against the original architectures that inspired components of BST, as well as modern standard Transformers.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The paper would be stronger with more discussion of limitations.

Summary

The authors present a novel architectural framework called the Block-State Transformer (BST), which integrates State Space models and Block-Recurrent Transformers to create a competitive autoregressive language model capable of effectively processing lengthy sequences. The input sequence is passed through a State Space model like S4, and the output of which is later in Block-recurrent Transformer as a replacement of the recurrent state vectors. To obtain the final output, the input embeddings are divided into fixed-sized windows and processed in parallel by a series of Block-Recurrent Transformers. Due to the usage of the S4 output as recurrent state vectors within the Block-Recurrent Transformers, the absence of recurrence allows for parallel computation. The authors propose three distinct integration modes, which differ in terms of how the S4 output is integrated within the recurrent state vectors. To evaluate the performance of BST, the authors compare it against four baseline models: Transformer-XL, Slide, Block-Recurrent Transformer, and a hybrid Gated State Space model. The comparison is conducted across three diverse datasets, namely PG19, arXiv Math, and GitHub. BST demonstrates slight perplexity improvements in the PG19 and GitHub datasets. Additionally, the authors present ablation studies on various parameters, including SSM layer placement, the number of SMM layers, and the state dimensionality of SMM.

Strengths

+ **Strong Presented Results**: Authors presented results on competitive benchmarks across reasonable prior baselines such as Transformer-XL, Slide, Block-Recurrent Transformers etc. and do outperform them across several tasks. + **Computational efficiency**: The proposed method is able to provide a huge improvement in terms of computational efficiency over models like Block-Recurrent Transformers by parallelization. + **Interesting combination of prior ideas**: By using the parallelizable nature of the SSMs the authors were able to introduce parallelization to Block-Recurrent Transformers thereby achieving computation efficiency.

Weaknesses

- **Incomplete Related works**: The authors' treatment of related works, particularly in the context of models combining Transformers and S4 appears to be lacking. It would have been beneficial for the authors to provide a more comprehensive discussion on existing models that incorporate both Transformers and S4. This would have allowed for a deeper exploration of the advancements and limitations of these models, highlighting the unique contributions of the proposed Block-State Transformer (BST). In addition to that, there could be more details on SSM development in related works since S4, S5, and S4D were mentioned in the paper later. - **Missing Preliminaries on S4/S5** : A more detailed description of the S4 models within the method section would have been beneficial, particularly regarding the computation of the kernel since the complexity of the S4 model is the major part of the BST model complexity. Specifically, the computation of the kernel is not trivial if you want to keep overall L log(L) complexity and it relies on the form of the A and B matrices. A more comprehensive exposition regarding the computational aspects of the S4 models is deemed necessary for a thorough understanding of the subject matter. - **Additional benchmarking** : As the authors themselves admit in the limitations section, there are further results required, especially on well benchmarked domains such as the Long Range Arena and also other long-term datasets to provide convincing evidence of BST performance.

Questions

Broadly, understanding BST capabilities in other settings would help enrich the paper results. See weaknesses.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

NA