Efficient Long Sequence Modeling via State Space Augmented Transformer

Weaknesses:

• We conduct additional experiments to demonstrate the effectiveness of SPADE. Specifically, we use S4, S5 [2], and Hyena [1] as the SSM in SPADE (Figure 3). We remark that other SSMs, such as GSS [3], are also outperformed by state-of-the-art models: S5 and Hyena. We also compare SPADE with positional embedding methods such as relative positional embedding and rotary positional embedding. We conduct language modeling experiments, and we report the ppl on the test set (lower the better). All models in the experiments have approximately 120M parameters.

• From the results, we can see that SPADE is easily extendable to accommodate various SSM. Moreover, for all configurations, we see that SPADE significantly outperforms the SSM-only model.

• In the experiments, we use the standard S4 as the SSM.

• We have included the code in the supplemental material.

[1] https://arxiv.org/pdf/2302.10866.pdf
[2] https://arxiv.org/pdf/2208.04933.pdf
[3] https://openreview.net/forum?id=5MkYIYCbva

Weaknesses:

• We conduct additional experiments to demonstrate the effectiveness of SPADE. Specifically, we use S4, S5 [2], and Hyena [1] as the SSM in SPADE (Figure 3). We remark that other SSMs, such as GSS [3], are also outperformed by state-of-the-art models: S5 and Hyena. We also compare SPADE with positional embedding methods such as relative positional embedding and rotary positional embedding. We conduct language modeling experiments, and we report the ppl on the test set (lower the better). All models in the experiments have approximately 120M parameters.

• From the results, we can see that SPADE is easily extendable to accommodate various SSM. Moreover, for all configurations, we see that SPADE significantly outperforms the SSM-only model.

Questions:

We explained the reason SSMs cannot handle local dependency well in Section 3.1. In language modeling, local information is more important than global information (in contrast to the Long Range Arena benchmark, which is specifically designed for long-range dependence modeling). This is because in practice, we rarely see words (tokens) that are thousands of positions apart exhibit strong dependencies [4]. Therefore, the results in Figure 2 demonstrates that SSMs cannot handle local information as well as attention models.

[1] https://arxiv.org/pdf/2302.10866.pdf
[2] https://arxiv.org/pdf/2208.04933.pdf
[3] https://openreview.net/forum?id=5MkYIYCbva
[4] https://arxiv.org/pdf/1905.07799.pdf

Weaknesses:

1. One of the properties of SSMs is that they do not need to be trained to capture global information (although training usually brings more gain). In our pre-training experiments, we do not train the SSM in the bottom layer, and we only train the attention and FFN weights. Therefore, even though we set the pre-training length to 512, our pre-trained model can handle any sequence length during fine-tuning, e.g., the summarization tasks we considered.

2. We conduct ablation experiments on the location and number of SSMs in Section 6.3. We conduct additional experiments to demonstrate the effectiveness of SPADE. Specifically, we use S4, S5 [2], and Hyena [1] as the SSM in SPADE (Figure 3). We remark that other SSMs, such as GSS [3], are also outperformed by state-of-the-art models: S5 and Hyena. We also compare SPADE with positional embedding methods such as relative positional embedding and rotary positional embedding. We conduct language modeling experiments, and we report the ppl on the test set (lower the better). All models in the experiments have approximately 120M parameters. From the results, we can see that SPADE is easily extendable to accommodate various SSM. Moreover, for all configurations, we see that SPADE significantly outperforms the SSM-only model.

Questions:

1. Our model is easier to scale compared with SSM-only models. In our experiments, we find that using a single SSM in the bottom layer suffices for the model to perform well. Therefore, the majority of SPADE are attention-based layers, which are easy to scale.

2. Thank you for the suggestion. We will further investigate how SSMs propagate global information. We would like to highlight that in SPADE, we use the SSM to model global information, and we use the local attention modules to capture more fine-grained local information,

[1] https://arxiv.org/pdf/2302.10866.pdf
[2] https://arxiv.org/pdf/2208.04933.pdf
[3] https://openreview.net/forum?id=5MkYIYCbva

Weaknesses:

1. The major contribution of the SPADE is its scalability and extensibility. Existing works on SSMs need significant engineering efforts to scale up. However, by combining SSMs with Transformer layers, the proposed method is easy to scale and can be easily extended to accommodate advanced SSM structures.

2. We address the following about baselines:

   • From the Hyena paper [1], RNN-like mechanisms such as RWKV have worse performance than SSM such as Hynea. We conduct additional experiments to demonstrate the effectiveness of SPADE. Specifically, we use S4, S5 [2], and Hyena [1] as the SSM in SPADE (Figure 3). We remark that other SSMs, such as GSS [3], are also outperformed by state-of-the-art models: S5 and Hyena. We also compare SPADE with positional embedding methods such as relative positional embedding and rotary positional embedding.We conduct language modeling experiments, and we report the ppl on the test set (lower the better). All models in the experiments have approximately 120M parameters.

   • From the results, we can see that SPADE is easily extendable to accommodate various SSM. Moreover, for all configurations, we see that SPADE significantly outperforms the SSM-only model.

Questions:

1. We pre-train a T5 model using the same setting as SPADE, and we include the fine-tuning results in Table 3. From the results, we see that the re-implemented T5 model outperforms the public T5 model. However, SPADE still significantly outperforms the re-implemented T5 model (86.8 vs 85.1 for average score on GLUE).

[1] https://arxiv.org/pdf/2302.10866.pdf
[2] https://arxiv.org/pdf/2208.04933.pdf
[3] https://openreview.net/forum?id=5MkYIYCbva

Weaknesses:

• One of the main contributions of SPADE is its scalability and extensibility. Existing works on SSMs need significant engineering efforts to scale up. However, by combining SSMs with Transformer layers, the proposed method is easy to scale and can be easily extended to accommodate advanced SSM structures.

• In Figure 6, we systematically study why the proposed architecture is effective. Specifically, through extensive experiments, we demonstrate that using a single global layer (i.e., a layer augmented with SSM) at the bottom of the model suffices for the model to perform well.

• We would like to highlight that existing models with SSMs focus on long-sequence modeling. However, through our experiments in Table 3, we demonstrate that SPADE is effective for both long-sequence and short-sequence modeling. This compensates a drawback of existing methods.

Questions:

• We also investigated other options: in Figure 6 we investigate the location of the SSM, and we conclude that the best location is to put the SSM at the bottom of the model. We also investigated other options to combine attention with LLM, such as stacking the SSM with attention (e.g., first compute the SSM, and then the output is fed to the attention). In practice, we only find marginal differences between this approach and the proposed method.

• For SSM such as S4 and Hyena, the complexity is $O(L \log (L))$, where $L$ is the sequence length. For a model with $N$ SSM layers, the total time complexity is $O(N L \log (L))$. In SPADE, we only use one SSM layer, and the other layers contain local attention modules with linear complexity. Therefore, the overall complexity of SPADE is $O(L \log (L) + (N-1) L)$, which is faster than models with only SSM layers. We demonstrate this in Figure 5.