Taipan: Efficient and Expressive State Space Language Models with Selective Attention

The proposed hybrid model is conceptually similar to other hybrid models that combine softmax attention models (Transformers) and modern CNN layers (such as S4, Mamba). Although the gain on small models is encouraging, the paper could be much stronger if a more comprehensive comparison with SOTA Transoformer / Hybrid models can be performed.

1. The comparison between Jamba and Taipan is not fair: Taipan uses 1:6 for number of attention layers v.s. number of Mamba layers, while Jamba uses 1:7. Also, Taipan uses Mamba 2 while Jamba uses Mamba 1. The performance gain of Taipan in Table 1 can be from the fact that Taipan uses Mamba 2 and has more attention layers, and may have nothing to do with the proposed selective attention.
2. Lack of novelty: Hybridization between state-space-models and dynamic selective SWA has been explored in SeqBoat [1], but the paper does not include any discussions or emipical study to compare different selction mechanisms. Also, Taipan does not select key-value pairs, which will limit its long context performance.
3. Lack of important baselines: The paper should at least compare the performance of Taipan with a simple baseline that has 1:6 SWA-Mamba2 ratio to prove the effectivnees of the proposed selective attention. A more thorough comparisons should includes different sparse attention baselines as proposed in BigBird [2] and LongFormer, which are now well supported by FlexAttention [3].
4. Lack of implementation details: The paper does not includes a detailed description of how hyperparameters are configurated, such as: the temperature of Gumbel softmax, and how the query selection is efficiently implemented so that the proposed selective attention can result in wall-time speed up.
5. Taipan only shows non-exploading perplexity for long context extrapolation, which is trival for SWA based Mamba hybrid models, considering that Samba [4] already shows improving perplexity up to 1M context length. The paper can be strengthened with more evidiences on long context tasks such as Passkey Retrieval.

Missing References:

[1] Sparse Modular Activation for Efficient Sequence Modeling (NeurIPS 2023)

[2] Big Bird: Transformers for Longer Sequences (NeuIPS 2020)

[3] https://github.com/pytorch-labs/attention-gym

[4] Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling (arXiv 2023)

1. Even though the goal of selective attention is to improve efficiency, the experimental section does not quantify the computational benefits in terms of memory and latency compared to full attention or different budgets in practice. I'd suggest extending the experiment in Figure 5 to include memory use and training/inference times.
2. The comparison to previous efficient and hybrid models has limited coverage as it included only two baseline models and model sizes up to 1.3B. This reduces the potential impact of the main findings. To strengthen the claims regarding scaling, I'd suggest adding a larger model to reach the 7B mark and including a table with results compared to other recent efficient or hybrid architectures such as RecurrentGemma.
3. The experiment scope could benefit from recent general evaluation benchmarks for LLMs (MMLU, HELM, BBH), and instruction tuning or preference optimization experiments, with higher priority on the general evaluation. The effect of different hyper-parameters such as sliding window size from 64 up to maximum context length in a controlled experiment would also be useful.

The concept of selective attention is promising, as it aligns well with recent advances in efficient language models. However, similar approaches have been explored in prior work, including "Power-BERT: Accelerating BERT Inference via Progressive Word-vector Elimination" and "A Gated Self-attention Memory Network for Answer Selection." These studies also leverage selective focus on important tokens, prioritizing computation for tokens requiring additional context. Further distinction from these works, especially in terms of innovation and unique contributions, would enhance the impact of this paper.

Among previous research, "Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling" appears most comparable due to its hybrid structure and sliding window mechanism in attention. I would anticipate that Samba could achieve similar results in performance and latency to the model in this paper. A thorough empirical comparison with Samba would be beneficial to underscore the advantages and trade-offs of the proposed approach.

In Figure 1, perplexity seems to increase steadily from the start. Typically, one might expect an initial decrease in perplexity with context length before a rise as the length extends beyond a certain threshold such as the pre-training context length. Additionally, the claim that Taipan outperforms Mamba in terms of latency is unclear. Providing further clarification on latency measurements and factors contributing to Taipan’s efficiency would enhance the reader’s understanding.

Regarding task performance, additional explanation is needed to clarify why Taipan outperforms Transfer on the tasks listed in Table 1, as many involve short-context scenarios. More supporting evidence to validate Taipan’s superiority on these tasks would strengthen the claims. Furthermore, including benchmarks on MMLU and GSM-8K, which require higher reasoning capabilities, would offer a more comprehensive assessment of the model's generalization and reasoning skills.

I believe the paper has several notable weaknesses that limit its impact.

Efficiency: First, the presentation of efficiency gains is potentially misleading in Figure 1b, as Taipan’s backbone, Mamba-2, is slower than Taipan itself. Either that line represents Mamba-1, or the plot should include Mamba-2. To make matters more confusing, line 428 states, "Notably, Taipan consistently outperforms Mamba-2, primarily due to its selective attention mechanism." Therefore, how is it possible for a model that uses Mamba-2 to process the input, along with additional computations, to actually be faster than Mamba-2? Overall, this discrepancy raises questions about whether computational overheads are fully accounted for.

Novelty: Furthermore, the novelty of combining SSMs with attention mechanisms is limited, as previous models, such as Jamba, have explored similar hybrid architectures, while the selective attention mechanism can be seen as an increment over Jamba.

Gumbel-softmax: Arbitrary architectural choices, like the selection of Gumbel-softmax without justification or comparison with alternatives, also weaken the paper, especially given that SALs are a primary contribution. The fixed attention capacity $C$ set during training could reduce the model’s flexibility at test time, as the need for attention across tokens may vary, and it is unclear how the model avoids processing all tokens at test time (as $C$ is budget for training).

Presentation: Additionally, inconsistencies in result reporting (e.g., bolded Taipan results even where it does not outperform other models) could mislead readers, as could unclear visual elements like Figure 1’s unexplained extrapolation regime and Figure 4’s table format. Moreover, the paper also disregards the proper use of citation styles (citep vs citet). Regarding Figure 1, it is unclear where the extrapolation regime starts, as per section 4.4. Collectively, these issues make the paper feel overly incremental and poorly substantiated. Therefore, to improve the paper, I suggest consistently bolding the best results, clearly highlighting the extrapolation regime in Figure 1, improving Figure 4, and fixing the citation format throughout the paper.