Paper: Efficient Long Sequence Modeling via State Space Augmented Transformer (ICLR-2024)

Reviewer comment (Experiments):
My main concern for this paper is the lack of baselines. Including a comparison with other recent SSMs, such as S5 [1] or Hyena [2], would be highly beneficial.

Author response:
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.

Reviewer comment (Experiments):
The abovementioned models could be incorporated with SPADE by replacing the SSM module with any other model. However, it would be beneficial to understand whether the performance gap between SPADE and S4 is caused by adding Transformer blocks on top of the SSM module. Training SPADE with Hyena may not improve performance over Hyena while being better than SPADE with S4. In this case, the motivation for the paper will disappear.

Author response:
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.

Reviewer comment (Writing):
I struggled to understand which SSM was used in SPADE (did I miss it?). It should be S4 based on the text. However, it would be helpful to name it in Section 3.2 explicitly.

Author response:
In the experiments, we use the standard S4 as the SSM.

Reviewer comment (Reproducibility):
The paper lacks reproducibility despite an extensive description of training within the text since no supplementary material with source code was released.

Author response:
We have included the code in the supplemental material.

Reviewer comment (Experiments):
In the experiments, the performance gain compared to previous (truncated) transformer approaches is quite limited. The proposed method is viewed as a novel position encoding mechanism, and a comparison against vanilla/truncated Transformers with more advanced position encodings, such as Rotary embeddings, ALiBi, and/or Transformer-XL, is expected.

Author response:
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.


|  | ppl |
|---|---|
| Transformer | 22.3 |
| Transformer-Relative | 20.9 |
| Transformer-Rotary | 20.8 |
| S4 | 22.7 |
| SPADE-S4 | 20.2 |
| S5 | 23.0 |
| SPADE-S5 | 20.0 |
| Hyena | 22.9 |
| SPADE-Hyena | 20.3 |

[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

Reviewer comment (Experiments):
There is confusion about the fact that while SSM alone cannot build a successful language model, using it (as a replacement of the position encoding) along with a truncated (windowed/chunked) Transformer results in an efficient long-term dependency capturing mechanism. An ablation study is requested to show that inferior versions of SSM, while still computationally efficient, are not capable enough to support the dependency-capturing capabilities in language modeling compared to the proposed SPADE.

Author response:
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.

[4] https://arxiv.org/pdf/1905.07799.pdf

Reviewer comment (Novelty):
S4 and Transformer are widely-used models, combining them together brings somehow incremental novelty contributions.

Author response:
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.

Reviewer comment (Experiments):
The comparisons with alternative methods for capturing global context, such as RNN-like mechanisms or other efficient attention mechanisms, are incomplete and lack essential details. Only S4 is compared in table 1. Some recent long-sequence modeling studies:  
a. LongNet: Scaling Transformers to 1,000,000,000 Tokens  
b. Long Range Language Modeling via Gated State Spaces

Author response:
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.

|  | ppl |
|---|---|
| Transformer | 22.3 |
| Transformer-Relative | 20.9 |
| Transformer-Rotary | 20.8 |
| S4 | 22.7 |
| SPADE-S4 | 20.2 |
| S5 | 23.0 |
| SPADE-S5 | 20.0 |
| Hyena | 22.9 |
| SPADE-Hyena | 20.3 |


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

Reviewer comment (Novelty):
The paper primarily consolidates existing concepts, such as SSM and local attention, and offers limited novelty in terms of new methodologies.

Author response:
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.

Reviewer comment (Theory):
Aside from the experimental results, it falls short in providing a comprehensive understanding of why this architecture is effective, and more crucially, in identifying scenarios where this approach may not be as effective.

Author response:
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.

Reviewer comment (Experiments):
In Figure 3, you simply concatenate SSM and local attention output, and apply a weight $\bf W$, did you try any other method to fuse them?

Author response:
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.