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

## Abstract

Abstract Efficient long-context language modeling remains a significant challenge in Natural Language Processing (NLP). While Transformers dominate language tasks, they struggle with long sequences due to quadratic computational complexity in training and linearly scaling memory costs during inference. Recent State Space Models (SSMs) such as Mamba offer alternatives with constant memory usage, but they underperform in tasks requiring extensive in-context retrieval. We introduce Taipan, a novel hybrid architecture that combines Mamba-2 with Selective Attention Layers (SALs). These SALs identify tokens requiring long-range interactions, remove less important features, and then augment their representations using the attention module. This approach balances Mamba’s efficiency with Transformer-like performance in memory-intensive tasks. By constraining the attention budget, Taipan extends accurate predictions to context lengths of up to 1 million tokens while preserving computational efficiency. Our experiments demonstrate Taipan’s superior performance across various scales and tasks, offering a promising solution for efficient long-context language modeling.

## 1 Introduction

Transformer-based architectures Vaswani (); Brown () have revolutionized Natural Language Processing (NLP), delivering exceptional performance across diverse language modeling tasks Touvron et al. (). This success stems from their ability to capture complex word dependencies using the self-attention mechanism. In addition, Transformers are highly scalable and well-suited for parallel training on large datasets. However, despite their success, they still face notable challenges when handling long-context sequences. Specifically, the self-attention mechanism suffers from quadratic computational complexity, and the memory requirement grows linearly with context length during inference, as the model must store key-value vectors for the entire context. These factors impose practical constraints on sequence length due to the high computational and memory costs.

To this end, recent advancements in recurrent-based architectures, particularly State Space Models (SSMs) Gu et al. (; ), have emerged as promising alternatives for efficient language modeling Gu & Dao (); Dao & Gu (). SSMs offer constant memory usage during inference, and architectures like Mamba-2 Dao & Gu (), a variant of SSMs, have demonstrated performance comparable to Transformers in certain language tasks Waleffe et al. (). Some studies even suggest that SSMs can outperform Transformers in areas like state tracking Merrill et al. () due to their Markovian nature. However, despite these advancements, SSM-based models still fall short in scenarios requiring in-context retrieval or handling complex long-range dependencies Arora et al. (); Waleffe et al. ().

To address these challenges, we introduce Taipan, a hybrid architecture that combines the efficiency of Mamba with enhanced long-range dependency handling through Selective Attention Layers (SALs). While Mamba is highly efficient, it relies on the Markov assumption—where predictions are based solely on the last hidden state—which can lead to information loss for tokens that need interactions with distant tokens. To mitigate this, Taipan incorporates SALs that strategically select key tokens in the input sequence requiring long-range dependencies. These selected tokens first undergo feature refinement to remove unimportant information, and then are passed through an attention module to capture long-range dependencies.

Less critical tokens bypass the attention step, as we hypothesize that their Markovian representations from Mamba contain sufficient information for accurate prediction, obviating the need for additional attention-based augmentation. This selective approach enables Taipan to balance Mamba’s computational efficiency with enhanced long-range modeling capabilities.

SALs play a crucial role in Taipan’s design, both in enhancing performance and ensuring computational efficiency. By focusing the attention mechanism on a subset of important tokens, SALs reduce the computational costs that come from attention modules. This targeted approach enables Taipan to excel in memory-intensive tasks while maintaining efficiency during both training and inference. Importantly, Taipan retains the linear memory usage characteristic of SSMs, offering a significant advantage over traditional Transformer models in handling extremely long sequences.

We scale Taipan to $190$M, $450$M, and $1.3$B parameters, pre-training on $100$B tokens. Experimental results demonstrate Taipan’s superior performance across a wide range of tasks. In zero-shot language modeling evaluations, Taipan consistently outperforms both Transformer and Mamba baselines, showcasing its strong general language understanding capabilities. More notably, in memory-intensive tasks such as long-context retrieval and structured information extraction, Taipan exhibits significant improvements over Mamba-2, addressing a key limitation of existing recurrent-based models. Furthermore, Taipan demonstrates remarkable extrapolation capabilities, maintaining high performance on sequences up to $1$ million tokens in context-length - while preserving efficient generation capabilities. This combination of broad task proficiency, superior performance in memory-intensive scenarios, and exceptional long-context modeling positions Taipan as a versatile and powerful architecture for advanced language processing tasks.

Figure: Figure 1: Model Performance Comparison. a) Perplexity across different context lengths. Lower perplexity indicates better performance. b) Latency comparison of models at various generation lengths. Taipan exhibits significantly lower latency and superior scaling compared to other strong baselines for longer sequences.
Refer to caption: x1.png

## 2 Background

This section briefly overviews the foundational architectures relevant to our work. We first review Causal Self-Attention Vaswani (), the core mechanism of Transformer models. We then discuss Linear Attention Katharopoulos et al. (), an efficient variant that achieves linear complexity. Finally, we examine Mamba-2, a recent architecture that generalizes Linear Attention using structured state-space models (SSMs) Dao & Gu (). We emphasize how each model balances computational efficiency and recall accuracy, particularly in memory-intensive tasks .

### 2.1 Causal Self-Attention

Causal Self-Attention is the key component in Transformer architectures that allows each token in a sequence to attend to all other previous tokens (Vaswani, ). Given an input sequence $\mathbf{X}=[\mathbf{x}_{1},\ldots,\mathbf{x}_{L}]\in\mathbb{R}^{L\times d}$, where $L$ is the sequence length and $d$ is the embedding dimension, self-attention firsts computes the query, key, and value vectors for each token via linear projections:

$$ $\displaystyle\mathbf{q}_{i}=\mathbf{W}_{Q}\mathbf{x}_{i},\quad\mathbf{k}_{i}= \mathbf{W}_{K}\mathbf{x}_{i},\quad\mathbf{v}_{i}=\mathbf{W}_{V}\mathbf{x}_{i}$ $$

where $\mathbf{W}_{Q},\mathbf{W}_{K},\mathbf{W}_{V}\in\mathbb{R}^{d\times d}$ are learnable weight matrices.

Then, the attention output $\mathbf{o}_{i}$ for each token $\mathbf{x}_{i}$ will be calculated as a weighted sum of the value vectors over the distribution of similarity matrix between its query vector and previous key vectors:

$$ $\displaystyle\mathbf{o}_{i}=\sum_{t=1}^{i}{\frac{\exp(\mathbf{q}_{i}^{\top} \mathbf{k}_{t}/\sqrt{d})}{\sum_{j=1}^{i}\exp(\mathbf{q}_{i}^{\top}\mathbf{k}_{ j}/\sqrt{d})}}\mathbf{v}_{t}$ $$

The non-linear softmax distribution allows the models to capture intricate relationships between tokens, and concentrate on salient features Qin et al. (); Zhao et al. (). As such, self-attention can encode complex language patterns and long-range dependencies that are crucial for complex language understanding and generation tasks.

### 2.2 Linear Attention

To address the quadratic complexity, recent work has shown that it is possible to achieve linear complexity with the attention mechanism by replacing the softmax attention with dot-product attention (Shen et al., ; Katharopoulos et al., ). Given a feature transformation $\phi(\mathbf{x})$, causal self-attention can be rewritten as:

$$ $\displaystyle\mathbf{o}_{i}=\sum_{t=1}^{i}{\frac{\phi(\mathbf{q}_{i})^{\top} \phi(\mathbf{k}_{t})}{\sum_{j=1}^{i}\phi(\mathbf{q}_{i})^{\top}\phi(\mathbf{k} _{j})}}\mathbf{v}_{t}$ $$

Then, using the associate property of matrix multiplication, this can be reformulated as:

$$ $\displaystyle\mathbf{o}_{i}=\frac{\phi(\mathbf{q}_{i})^{\top}\sum_{t=1}^{i} \phi(\mathbf{k}_{t})\mathbf{v}_{t}^{\top}}{\phi(\mathbf{q}_{i})^{\top}\sum_{t= 1}^{i}\phi(\mathbf{k}_{t})}$ $$

Let $\mathbf{S}_{i}=\sum_{t=1}^{i}\phi(\mathbf{k}_{t})\mathbf{v}_{t}^{\top}$ and $\mathbf{z}_{i}=\sum_{t=1}^{i}\phi(\mathbf{k}_{t})$. We can then rewrite the equation in a recurrent form:

$$ $\displaystyle\mathbf{S}_{i}$ $\displaystyle=\mathbf{S}_{i-1}+\phi(\mathbf{k}_{i})\mathbf{v}_{i}^{\top}$ $\displaystyle\mathbf{o}_{i}$ $\displaystyle=\frac{\mathbf{S}_{i}\phi(\mathbf{q}_{i})}{\mathbf{z}_{i}^{\top} \phi(\mathbf{q}_{i})}\approx\mathbf{S}_{i}\phi(\mathbf{q}_{i})$ $$

This formulation allows for efficient training and inference. Let $\mathbf{Q},\mathbf{K},\mathbf{V}\in\mathbb{R}^{L\times d}$ be the query, key, and value matrices of the sequence input $\mathbf{X}$.
During training, we can use the matrix multiplication form: $\mathbf{O}=(\mathbf{Q}\mathbf{K}^{\top}\odot\mathbf{M}_{L})\mathbf{V}$, where $\mathbf{M}_{L}$ is a causal mask. At inference time, we can use the recurrent form for efficient sequential processing.

However, despite its computational efficiency, linear attention has notable limitations compared to softmax attention. The dot-product approximation in linear attention lacks the nonlinear normalization of softmax, often resulting in a more uniform distribution of attention weights Han et al. (). This uniformity can impair the model’s ability to focus sharply on specific and relevant tokens. Consequently, linear attention models may underperform in tasks requiring precise in-context retrieval or focused attention on particular input segments Han et al. ().

### 2.3 Mamba-2

Mamba Gu & Dao () is a variant of structured state space models (SSMs) that uses the selective data-dependent mechanism. Mamba-2 Dao & Gu () builds on this foundation, revealing deep connections between SSMs and linear attention Katharopoulos et al. () through the framework of structured state-space duality (SSD).

The core of Mamba-2 can be defined by using the recurrent form:

$$ $\displaystyle\mathbf{h}_{t}$ $\displaystyle=\mathbf{A}_{t}\mathbf{h}_{t-1}+\mathbf{B}_{t}\mathbf{x}_{t}$ $\displaystyle\mathbf{o}_{t}$ $\displaystyle=\mathbf{C}_{t}\mathbf{h}_{t}$ $$

where $\mathbf{A}_{t}$ is further simplified to a scalar multiplied by the identity matrix. This formulation allows Mamba-2 to be interpreted as a generalization of linear attention.

The key insight of Mamba-2 is that this recurrence can be equivalently expressed as a matrix multiplication:

$$ $\displaystyle\mathbf{O}_{t}=(\mathbf{L}_{t}\odot\mathbf{C}_{t}\mathbf{B}_{t}^{ \top})\mathbf{X}_{t}$ $$

where $\mathbf{L}$ is a 1-semiseparable matrix. This matrix form reveals the duality between the recurrent (linear-time) and attention-like (quadratic-time) computations. Also, the 1-semiseparable matrix $\mathbf{L}$ encodes the temporal dependencies, while $\mathbf{C}\mathbf{B}^{\top}$ represents content-based interactions similar to attention. This formulation generalizes linear attention, which can be seen as a special case where $\mathbf{L}$ is the all-ones lower triangular matrix.

While Mamba-2 is efficient, it shares the same limitations as Linear Attention in terms of precise memory recall Arora et al. (); Wen et al. (), leading to reduced performance in tasks that demand accurate retrieval of specific sections in the input sequence.

## 3 Taipan Model

Figure: Figure 2: An overview of the Taipan architecture.
Refer to caption: x2.png

To address the limited modeling capabilities of Mamba-2 and Linear Attention while preserving their computational efficiency, we introduce Taipan, a new architecture for sequence encoding in language modeling. In Taipan, we strategically incorporate Selective Attention Layers (SALs) within the Mamba framework, as shown in Figure 2. SALs are inserted after every $K$ Mamba-2 blocks, creating a hybrid structure that combines Mamba-2’s efficiency with Transformer-style attention for effective sequence representation.

The core of SALs is a gating network that identifies important tokens for enhanced representation modeling. These tokens undergo two phases: (1) feature refinement to filter out irrelevant information and (2) representation augmentation via softmax attention. This allows Taipan to capture complex, non-Markovian dependencies when necessary.

Taipan processes input through Mamba-2 blocks, with SALs periodically refining key token representations. These enhanced representations are then passed into the subsequent Mamba-2 layers, influencing further processing. This hybrid structure balances Mamba-2’s efficiency with the expressive power of SALs, enabling Taipan to excel in tasks requiring both speed and accurate information retrieval. The following sections detail each component’s structure and function.

### 3.1 Selective Attention Layers

Figure: Figure 3: Attention mechanisms in Taipan’s Selective Attention Layers. White areas indicate no attention. (a) Full Causal Attention (b) Sliding Window Attention ($w=4$) (c) Selective Attention ($C=0.3$, $w=5$)
Refer to caption: x3.png

Selective Attention Layers (SALs) are the key innovation in Taipan, designed to enhance the model’s ability to focus on critical tokens while maintaining overall efficiency. These layers employ a lightweight gating network $G_{\theta}$ to dynamically determine which tokens should undergo softmax attention processing.

For each token hidden representation $\mathbf{h}_{i}$ in the input sequence, the gating network $G$ computes a score vector:

$$ $\mathbf{s}_{i}=G_{\theta}(\mathbf{h}_{i})$ (1) $$

where $G_{\theta}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{2}$ is parameterized by $\theta$. This score vector $\mathbf{s}_{i}=[s_{i,0},s_{i,1}]$ serves two purposes: 1) it is used to generate a binary mask $m_{i}$ for token selection, and 2) it guides feature refinement.

To maintain differentiability while allowing for discrete token selection, we employ the Straight-Through Gumbel-Softmax trick Jang et al. (). A binary mask $m_{i}$ is generated from $\mathbf{s}_{i}$ to select tokens during the forward pass of the network:

$$ $m_{i}=\text{argmax}(\text{GumbelSoftmax}(\mathbf{s}_{i},\tau))$ (2) $$

where $\tau$ is the temperature parameter. $\mathbf{h}_{i}$ will only be selected for attention processing if $m_{i}=1$.

For the backward pass, we instead use continuous Gumbel-Softmax approximation of $m_{i}$ to achieve computation differentiability for the network:

$$ $\tilde{m}_{i}=\frac{\mathbb{I}[m_{i}=0]\exp((\mathbf{s}_{i,0}+g_{0})/\tau)+ \mathbb{I}[m_{i}=1]\exp((\mathbf{s}_{i,1}+g_{1})/\tau)}{\exp((\mathbf{s}_{i,0} +g_{0})/\tau)+\exp((\mathbf{s}_{i,1}+g_{1})/\tau)}$ (3) $$

where $\mathbb{I}[\textperiodcentered]$ is the indicator function, and $g_{0}$ and $g_{1}$ are i.i.d samples from the Gumbel$(0,1)$ distribution. In this way, we are able to train our entire model, including the gating network, in an end-to-end fashion for language modeling.

For the selected tokens (those with a mask value $m_{i}$ of 1), we compute their attention-based representations:

$$ $\mathbf{o}_{i}=\text{Attention}(\mathbf{q}_{i},\mathbf{K},\mathbf{V})$ (4) $$

where $\mathbf{q}_{i}$ is the query vector for the $i$-th selected token (denoted $\mathbf{h}_{i}^{s}$), and $\mathbf{K}$ and $\mathbf{V}$ are the key and value matrices for previous tokens.

In our model, the score vector $\mathbf{s}_{i}$ is also used to refine the representations of selected tokens. We employ the softmax of $\mathbf{s}_{i}$ to compute the mixing weights:
$[1-\alpha_{i},\alpha_{i}]=\text{softmax}(\mathbf{s}_{i})$.
The final output for a selected token $\mathbf{h}_{i}^{s}$ is a weighted combination:

$$ $\mathbf{h}_{i}^{s}=(1-\alpha_{i})\mathbf{h}_{i}^{s}+\alpha_{i}\mathbf{o}_{i}$ (5) $$

As such, Taipan can adaptively preserve key information in $\mathbf{h}_{i}^{s}$ while enriching the representation with the attention output $\mathbf{o}_{i}$.
In other words, $\alpha_{i}$ acts as the data-dependent factor, filtering out unimportant features from the original representation while integrating richer information from the attention outputs. Here, it is important to note that unselected tokens (i.e., $m_{i}=0$) skip the attention module and retain their original representations from Mamba-2.

Finally, all token representations are passed through a residual SwiGLU Shazeer () layer:

$$ $\mathbf{h}=\mathbf{h}+\text{SwiGLU}(\mathbf{h})$ (6) $$

This final transformation ensures that all token representations undergo consistent non-linear processing before being passed to the next layer in the network, enhancing the model’s ability to capture complex dependencies.

### 3.2 Sliding Window Attention

To maintain linear time complexity while leveraging the benefits of attention, Taipan employs Sliding Window Attention (SWA) Beltagy et al. (). SWA’s computational complexity scales linearly with sequence length, allowing Taipan to handle theoretically unlimited context lengths during inference. Importantly, the combination of Selective Attention and Sliding Window Attention in Taipan leads to a significantly sparser attention weight map compared to full attention or standard windowed attention (Figure 3), thus enhancing the computational efficiency of Selective Attention for processing long sequences for our model. In addition, the sparser attention map allows us to afford a longer sliding window (i.e., $w=2048$ in our work) to effectively capture longer-range dependencies for input sequences. In this way, our designed Taipan architecture offers a mechanism to balance the efficient processing of long sequences with the ability to capture important long-range dependencies, thereby addressing a key limitation of existing efficient attention mechanisms. Finally, removing positional embeddings from the Attention Module improves extrapolation capabilities, suggesting that the model can better generalize temporal relationships. We explore this impact of positional embeddings in more detail in Section 5.2.

### 3.3 Training and Inference

To better balance efficiency and expressiveness, we introduce an attention budget constraint. Given a predefined budget $C$, representing the desired fraction of tokens to receive attention, we incorporate a constraint loss into our training objective:

$$ $\displaystyle\mathcal{L}_{\text{constraint}}=\sum_{n=1}^{N}\left\|C-\frac{\sum _{i=1}^{L}m_{i}}{L}\right\|_{2}^{2}$ (7) $$

Here, $N$ is the number of SALs, $L$ is the sequence length, and $\sum_{i=1}^{L}m_{i}$ represents the number of tokens selected for attention processing. During training, we employ the Straight Through Gumbel Softmax estimator for $\tilde{m}_{i}$ in the backward pass Jang et al. (); Bengio et al. (), ensuring differentiability while maintaining discrete token selection in the forward pass, thereby enabling end-to-end training of the entire model. As such, our overall training objective includes a standard cross-entropy loss $\mathcal{L}_{\text{CE}}$ for language modeling and the budget constraint term: $\mathcal{L}=\mathcal{L}_{\text{CE}}+\lambda\mathcal{L}_{\text{constraint}}$, where $\lambda$ is a hyperparameter.

During inference, Taipan processes input tokens sequentially through Mamba-2 blocks. At each Selective Attention Layer, the gating network $G_{\theta}$ computes a score vector $\mathbf{s}_{i}=G_{\theta}(\mathbf{h}_{i})$ for each token representation $\mathbf{h}_{i}$. This score computes a binary mask $m_{i}$ to determine if $\mathbf{h}_{i}$ should be used for attention processing. Consequently, our selective attention approach maintains Mamba-2’s efficiency for most tokens while applying targeted attention to critical elements, enabling effective long-range dependency modeling with minimal computational overhead.

## 4 Experiments

We conducted extensive experiments to evaluate Taipan’s performance across various scales and tasks. Our evaluation strategy focuses on three main areas: (1) zero-shot evaluation on diverse benchmarks to demonstrate Taipan’s general language understanding capabilities (Section 4.2), (2) in-context retrieval tasks to assess Taipan’s ability to retrieve information from historical contexts (Section 4.3), and (3) extrapolation ability in long-context scenarios to evaluate performance on extremely long sequences (Section 4.4).

### 4.1 Experimental Setup

We evaluate Taipan across three model sizes: $190$M, $450$M, and $1.3$B parameters. To ensure a comprehensive and fair comparison, we benchmark Taipan against three strong baselines:

- •
Transformer++ Touvron et al. (): An enhanced version of the LLaMA architecture Touvron et al. (), incorporating Rotary Positional Embeddings Su et al. (), SwiGLU Shazeer (), and RMSNorm Zhang & Sennrich ().
- •
Mamba-2 Dao & Gu (): A state-of-the-art linear RNN model based on State Space Models (SSMs). Each Mamba-2 block consists of a depthwise convolutional layer Poli et al. (); Gu & Dao (), an SSM layer Dao & Gu (), and MLP layers.
- •
Jamba Lieber et al. (): A hybrid model combining full Causal Self-Attention layers (with Rotary Position Embedding Su et al. ()) and Mamba-2 layers. Unlike Taipan, Jamba uses full Causal self-attention instead of selective attention, retains positional embeddings, and lacks a feature refinement mechanism.

Implementation Details
We train all models from scratch in three configurations: $190$M, $450$M, and $1.3$B parameters. The training process is consistent across configurations with the following hyperparameters: a batch size of $0.5$M tokens per step, a cosine learning rate schedule with $2000$ warm-up steps, and AdamW Loshchilov () optimization with a peak learning rate of $5e-4$ decaying to a final rate of $1e-5$. We apply a weight decay of $0.01$ and use gradient clipping with a maximum value of $1.0$. All models are trained with a fixed context length of $4096$ tokens.

The training data size varies by model scale: the 190M model is trained on $27$ billion tokens, while the $450$M and $1.3$B models are trained on $100$ billion tokens. The dataset details can be found in Appendix A.

For Taipan-specific implementation, we use a hybrid ratio of $6:1$, inserting a Selective Attention Layer (SAL) after every $6$ Mamba-2 Blocks. The Mamba-2 blocks are kept identical to the original work Dao & Gu (). We set the attention capacity $C=0.15$. The sliding window attention mechanism uses a window size ($w$) of $2048$ tokens.

**Table 1: Zero shot results of Taipan against baseline models.**
| Params | Model | Wino. | PIQA | Hella. | ARCE | ARCC | OB. | Truth. | RACE | BoolQ | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| & Data |  |  |  |  |  |  |  |  |  |  |  |
|  | Transformer++ | 47.1 | 60.9 | 27.9 | 42.2 | 20.5 | 18.9 | 42.9 | 25.4 | 57.2 | 38.1 |
| 190M | Mamba | 49.6 | 60.7 | 29.3 | 45.3 | 21.8 | 20.6 | 40.8 | 27.2 | 59.3 | 39.4 |
| 27B | Jamba | 49.9 | 60.3 | 29.2 | 46.3 | 21.4 | 18.5 | 39.8 | 27.4 | 58.6 | 39.1 |
|  | Taipan | 51.0 | 62.6 | 29.4 | 46.7 | 20.7 | 21.8 | 41.1 | 26.6 | 58.7 | 39.9 |
|  | Transformer++ | 51.5 | 67.6 | 42.3 | 60.8 | 27.7 | 33.4 | 39.2 | 30.5 | 54.7 | 45.3 |
| 450M | Mamba | 52.7 | 68.9 | 42.7 | 61.4 | 27.1 | 34.0 | 38.5 | 29.3 | 53.2 | 45.3 |
| 100B | Jamba | 53.1 | 69.3 | 44.3 | 62.6 | 28.7 | 34.4 | 37.5 | 31.3 | 55.7 | 46.3 |
|  | Taipan | 53.0 | 69.6 | 46.6 | 65.6 | 32.9 | 36.6 | 38.6 | 30.7 | 60.4 | 48.2 |
|  | Transformer++ | 53.8 | 71.6 | 53.8 | 63.2 | 36.3 | 36.4 | 44.0 | 31.2 | 59.4 | 49.9 |
| 1.3B | Mamba | 55.2 | 73.0 | 55.6 | 70.7 | 38.0 | 39.0 | 39.9 | 32.0 | 61.8 | 51.7 |
| 100B | Jamba | 54.7 | 73.8 | 55.8 | 69.7 | 37.6 | 41.8 | 40.4 | 32.8 | 59.2 | 51.8 |
|  | Taipan | 57.0 | 74.9 | 57.9 | 71.2 | 39.3 | 40.4 | 43.0 | 34.4 | 61.5 | 53.3 |

### 4.2 Language Modeling Performance

We report the zero-shot performance of Taipan and baseline models on a diverse set of commonsense reasoning and question-answering tasks. These include Winograd (Wino.) Sakaguchi et al. (), PIQA Bisk et al. (), HellaSwag (Hella.) Zellers et al. (), ARC-easy and ARC-challenge (ARCe & ARCc) Clark et al. (), OpenbookQA (OB.) Mihaylov et al. (), TruthfulQA (Truth.) Lin et al. (), RACE Lai et al. (), and BoolQ Clark et al. (). It is worth noting that these tasks are brief and do not involve in-context learning, thus inadequately demonstrating long-context modeling or in-context learning retrieval abilities.

Table 1 presents the zero-shot results for models of three sizes: $190$M, $450$M, and $1.3$B parameters. The results are evaluated using the lm-evaluation-harness(^1^11[https://github.com/EleutherAI/lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness)) Gao et al. () framework.

As can be seen, Taipan consistently outperforms the baseline models across most tasks for all model sizes. Notably, the performance gap widens as the model size increases, with the 1.3B Taipan model showing significant improvements over other baselines. This suggests that Taipan’s architecture effectively captures and utilizes linguistic patterns, even in tasks that do not fully showcase its long-context modeling capabilities.

### 4.3 In-Context Recall-Intensive Performance

To evaluate Taipan’s proficiency in precise in-context retrieval, we assessed all models on a set of recall-intensive tasks Arora et al. (). These tasks are designed to test a model’s ability to extract and utilize information from longer contexts, a capability particularly relevant to Taipan’s architecture. Our evaluation suite includes two types of tasks: structured information extraction and question answering. For structured information extraction, we used the SWDE and FDA tasks Arora et al. (), which involve extracting structured data from HTML and PDF documents, respectively. To assess question-answering capabilities, we employed SQuAD Rajpurkar et al. (), which requires models to ground their answers in provided documents.

Figure: Figure 4: Performance on in-context retrieval tasks.

Table 4 demonstrates Taipan’s significant performance advantages over both Mamba and Jamba in in-context retrieval tasks. Notably, Taipan achieves this superiority while consuming fewer computational resources than Jamba, which utilizes full attention mechanisms. This efficiency is attributed to Taipan’s architecture, which combines Mamba-like elements with selective attention mechanisms, allowing it to filter out less important features. We also notice that Transformers excel at memory-intensive tasks in this experiment; however, they are constrained by linear memory scaling with sequence length, limiting their effectiveness and applicability for very long sequences. In contrast, Taipan maintains constant memory usage, offering a more efficient solution for processing long documents.

### 4.4 Long-Context Extrapolation

Figure 1 illustrates Taipan’s superior performance in handling extended sequences compared to Transformer, Jamba, and Mamba models. In perplexity evaluations across context lengths from 1K to 1M tokens (Figure 1a), Taipan yields the lowest perplexity, particularly excelling beyond the training context length.

This performance contrasts sharply with other models: Transformers struggle with longer contexts due to quadratic computational complexity and linear memory scaling with sequence length, often leading to out-of-memory errors. Jamba, despite its hybrid nature, faces similar challenges due to its use of full attention mechanisms. Both Transformer and Jamba models exhibit limited extrapolation ability beyond their training context lengths. Mamba, while more efficient than Transformers and Jamba, still shows performance degradation for very long sequences.

Latency comparisons (Figure 1b) further highlight Taipan’s exceptional efficiency. It demonstrates the lowest latency among all models, with linear scaling across sequence lengths. This contrasts with the quadratic scaling of Transformers and higher latency growth of Jamba. Notably, Taipan consistently outperforms Mamba-2, primarily due to its selective attention mechanism.

## 5 Ablation Study

We conducted a comprehensive ablation study to investigate the effect of the two key components in Taipan’s architecture, i.e., the attention budget capacity $C$ and the inclusion of Positional Embeddings in the SALs, on its performance and efficacy.

### 5.1 Effect of Attention Budget Capacity

Our first experiment aimed to determine the optimal value of Capacity $C$ that would maintain computational efficiency while maximizing performance on downstream tasks. We trained multiple variants of Taipan, each with 1.3B parameters, using different Capacity $C$ values: $0.10$, $0.15$, $0.20$, and $0.25$. Each variant was trained for $24,000$ steps, allowing us to observe both the immediate impact of different $C$ values and their effect on model performance over time.

We evaluated the performance of each variant at regular intervals on two representative tasks: SWDE Arora et al. () (for structured information extraction) and HellaSwag Zellers et al. () (for commonsense reasoning). These tasks were chosen to assess both the model’s ability to handle long-context retrieval and its general language understanding capabilities.

Figure: Figure 5: Effect of Attention Budget Capacity $C$ on Taipan’s Performance
Refer to caption: x4.png

As illustrated in Figure 5, Taipan achieves optimal performance with a Capacity $C=0.15$. We observed that increasing $C$ beyond 0.15 does not lead to significant improvements in results while increasing computational costs. Conversely, reducing $C$ below $0.15$ resulted in a noticeable drop in performance on tasks requiring precise in-context retrieval or complex long-range dependencies. These findings support our hypothesis that computational demands vary across tokens, with many adequately represented by Mamba’s Markovian structure without requiring attention mechanisms.

By selectively applying attention only to tokens that benefit from it, Taipan optimizes resource allocation, enabling high performance while improving computational efficiency.

### 5.2 Impact of Positional Embeddings

Our second experiment investigated the impact of Positional Embeddings in Taipan’s Attention mechanism, focusing on the model’s ability to handle and generalize to various context lengths. We trained two variants of the $1.3$B parameter Taipan model for $24,000$ steps with a fixed context length of $4096$ tokens. One variant incorporates Rotary Positional Embeddings Su et al. () in the Selective Attention layers, while the other excludes them. Figure 6 illustrates the performance of both variants in terms of perplexity across different context lengths.

The results reveal that Taipan without Positional Embeddings performs superiorly in generalizing context lengths beyond the training context. Both variants show comparable performance for sequences similar to or shorter than the training context length. However, as the sequence length increases, the performance gap between the two variants widens, with Taipan without Positional Embeddings maintaining lower perplexity scores. This suggests that the absence of Positional Embeddings enables more robust scaling to longer sequences. We attribute this improved generalization to the model’s increased reliance on attention representation rather than positional biases.

Figure: Figure 6: Perplexity comparison of Taipan variants with and without Positional Embeddings across different context lengths. Lower perplexity indicates better performance.
Refer to caption: x5.png

## 6 Related Work

Our approach builds on a foundation of relevant previous research. We will now discuss key studies that inform our methodology.

State Space Models: SSMs have emerged as a promising approach in attention-free architectures for language processing tasks. These models offer improved computational and memory efficiency compared to traditional attention-based models. The development of SSMs has progressed through several key iterations: S4 Gu et al. () introduced the first structured SSM, focusing on diagonal and diagonal plus low-rank (DPLR) structures. Subsequent variants like DSS Gupta et al. (), S4D Gu et al. (), and S5 Smith et al. () improved on this foundation. Frameworks like GSS Mehta et al. (), H3 Fu et al. (), and RetNet Sun et al. () incorporated SSMs into broader neural network architectures, often combining them with gating mechanisms or efficient attention approximations. Recently, Mamba Gu & Dao () introduced time-varying or selective SSMs, which addresses limitations of static dynamics in previous SSMs by incorporating input-dependent state transitions, leading to improved performance in various tasks.

Hybrid Architecture: Several recent studies H3 Fu et al. (), Griffin De et al. (), Zamba Glorioso et al. (), Jamba Lieber et al. () suggest the potential of blending SSM and the attention mechanism. These hybrid designs show promise in outperforming both traditional Transformers and pure SSM architectures, such as Mamba, particularly in scenarios requiring in-context learning capabilities.

Long Context Models: Recent advancements in sequence modeling have pushed the boundaries of context length, each with distinct approaches and challenges. Recurrent Memory Transformer Bulatov et al. () demonstrated 1M token processing, but primarily on synthetic memorization tasks. LongNet Ding et al. () proposed scalability to 1B tokens, yet practical evaluations were limited to sequences under 100K tokens. Hyena/HyenaDNA Poli et al. (); Nguyen et al. () claimed 1M token context, but faced efficiency issues at longer lengths.
Mamba Gu & Dao () showed consistent improvements up to 1M tokens in DNA modeling and competitive performance across various language tasks.

## 7 Conclusion

Taipan presents a significant advancement in long-context language modeling by combining the efficiency of Mamba with strategically placed Selective Attention Layers. Our experiments demonstrate Taipan’s superior performance across various scales and tasks, particularly in scenarios requiring extensive in-context retrieval, while maintaining computational efficiency. A key insight is that not all tokens require the same computational resources. Taipan’s architecture leverages this observation through its selective attention mechanism, which dynamically allocates computational resources based on token importance. This hybrid approach addresses limitations of both Transformers and SSMs, offering a promising solution for efficient, large-scale language processing. Future work could explore further optimizations and applications of this architecture.

## Appendix A Datasets

Our training data comprises a diverse set of datasets, carefully curated to ensure breadth and depth across various domains. This diverse collection includes specialized mathematics datasets (MetaMathQA Yu et al. (), NuminaMath-CoT LI et al. (), OpenWebMath Paster et al. (), Orca-Math Mitra et al. ()), high-quality web data (Fineweb-Edu-dedup Ben Allal et al. ()), synthetic data (Cosmopedia-v2 Ben Allal et al. ()), code data (Starcoderdata-python-edu Li et al. ()), and general knowledge sources (Wikipedia). This comprehensive approach aims to enable our model to handle a wide array of language modeling tasks. The inclusion of both domain-specific and broad-coverage datasets is designed to enhance the model’s versatility and robustness across language modeling tasks.

All datasets were tokenized using the LLama3’s tokenizer Dubey et al. (), resulting in 300B tokens.

The training data size varies by model scale: the 190M model is trained on 27 billion tokens (exclusively from Cosmopedia-v2), while the 450M and 1.3B models are trained on 100 billion tokens sampled from the combination of datasets mentioned above. Below are detailed descriptions of each dataset used:

- 1.
MetaMathQA Yu et al. (): A comprehensive mathematics dataset designed to enhance the model’s mathematical reasoning and problem-solving abilities.
- 2.
NuminaMath-CoT LI et al. (): A chain-of-thought mathematics dataset that promotes step-by-step reasoning in mathematical problem-solving.
- 3.
Cosmopedia-v2 Ben Allal et al. (): A large-scale synthetic dataset for pre-training, consisting of over 39 million textbooks, blog posts, and stories.
- 4.
Fineweb-Edu-dedup Ben Allal et al. (): A high-quality subset of the FineWeb-Edu dataset, containing 220 billion tokens of educational web pages. This dataset was filtered using an educational quality classifier to retain only the most valuable educational content.
- 5.
OpenWebMath Paster et al. (): A diverse collection of mathematical content from over 130,000 different domains, including forums, educational pages, and blogs. It covers mathematics, physics, statistics, computer science, and related fields.
- 6.
Starcoderdata-Edu Li et al. (): A subset of the Starcoder dataset, specifically filtered for high-quality educational content related to Python programming. This dataset aims to enhance the model’s coding capabilities.
- 7.
Orca-Math Mitra et al. (): A dataset focused on mathematical word problems, designed to improve the model’s ability to interpret and solve practical mathematical scenarios.
- 8.
Wikipedia: An English Wikipedia dataset providing a broad range of general knowledge across various topics.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: - Can you clarify which version of Mamba was used in Figure 1b? Can you provide a more detailed breakdown of the computational costs for both Taipan and Mamba-2?

- How does Taipan avoid the risk of all tokens being passed to the attention module at test time if the fixed attention capacity C is exceeded?

- Can you provide a rationale for choosing Gumbel-softmax? What would be potential alternatives? For example, how does Taipan compare with other differentiable strategies, such as gradient surrogates, continuous relaxations, etc, which have been shown to be effective in similar applications? See [1] for a comprehensive overview.

[1] Discrete Latent Structure in Neural Networks (https://arxiv.org/abs/2301.07473)
soundness: 2
strengths: I enjoyed the effort put in this paper towards balancing Mamba-2's efficiency with selective attention mechanisms, an approach that can offer benefits for handling long contexts. I also liked that the experimental setup contains multiple evaluations across various benchmarks and model scales, allowing some insight into Taipan’s potential in extended context scenarios.
summary: The paper presents Taipan, a hybrid model that incorporates attention modules into  Mamba-2 (an SSM). Specifically, it proposes to use Selective Attention Layers (SALs) to manage long-context tasks more efficiently in language modeling, such that only selected tokens are passed to (windowed) attention modules. In that way, Taipan selectively attends to critical tokens within an input, allowing it to capture long-range dependencies while also seeking to maintain computational efficiency.
weaknesses: 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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. What is the computational benefit for different attention budgets and usage in different layers compared to full attention in terms of memory and latency? It would be useful if the authors provide some empirical evidence to quantify the benefits of SALs. 
2. Could the authors include a few additional baselines in the datasets under study? I'd suggest to report scores in a table from prior work with efficient or hybrid architectures on the same datasets (e..g RecurrentGemma).
 3. The comparison to a baseline that uses only sliding window attention with the same window size as Tapain is missing. Could the authors report scores for this baseline across the tasks used in the experiment section? This would help to better understand the impact of selective attention.
soundness: 3
strengths: 1. Existing hybrid architectures that combine state-space models with standard attention typically are applied to a small subset of layers for all tokens. The idea of applying attention only to a specific subset of positions dynamically through selective attention is novel and provides a more efficient way to augment recurrent models with attention.  
2. The proposed model outperforms consistently two strong baselines that represent efficient and hybrid models, namely Mamba and Jamba, on general and recall-intensive tasks. At the same time, it outperforms standard attention models represented by Transformer++ or is behind by a moderate margin (~10% relative in recall-intensive tasks). 
3. Apart from the quality, the proposed model has superior extrapolation capabilities up to 1M tokens and lower latency with increasing context size compared to the aforementioned baselines.
summary: This paper proposes a new hybrid architecture that combines the recurrent formulation of state-space models with selective attention layers (SALs). The key component introduced, SAL, identifies tokens that have long-context dependencies, refines their features and augments their representations with standard attention. This aims to increase performance on tasks that require long-context memory without incurring the quadratic cost in standard Transformers. The evaluation performed on various tasks  and extrapolation shows superior performance compared to standard (Transformer++), efficient (Mamba) and hybrid (Jamba) models.  In addition, performance on recall-intensive tasks and model sizes up to 1.3B are also encouraging.
weaknesses: 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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: see weaknesses
soundness: 3
strengths: 1. The paper’s motivation is clear, focusing on an important and timely topic with practical significance.
2. The writing is clear and well-organized, making it easy to understand.
3. The selective attention concept is well-founded and adds a valuable perspective to the field
summary: This paper presents a hybrid approach that combines elements of Mamba and Transformer architectures, aiming to address two major challenges: the high computational complexity of Transformers in handling long contexts and the quality degradation issues encountered with Mamba. This approach aligns with prior research, including methods like Samba and Jamba.

The key contribution of the paper is its selective mechanism for token attention calculation. By incorporating a gating network, the model selectively skips attention computation for certain tokens, reducing inference costs. This optimization enhances the efficiency of the attention layer, achieving a notable speed-up without compromising performance and the paper demonstrates this with empirical results.
weaknesses: 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.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Line 337: Jamba does not have positional embedding.

2. How is the model performance sensitive to the $\lambda$ of the constraint loss?
soundness: 1
strengths: 1. The proposed approach is simple and shows more stable extrapolation on perplexity compared to Jamba.
2. The feature refinement mechanism, which uses the underlying probability distributino of the gumbel softmax to interporlate the residual branch output and the layer input, looks interesting.
summary: The paper proposes to sparsify the query of self-attention layers in the context of layerwise hybridization between state-space models and self-attention.
weaknesses: 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)

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: It is useful to run ablation experiments to show (1) the gain of using sliding window attention; (2) the gain over using only sliding window attention without selective attention; and (3) the number of tokens selected based on Eq. (1) and (2).

It is also useful to investigate what tokens are selected, and whether there are any patterns, such as the ones described in https://arxiv.org/pdf/2310.01801.
soundness: 3
strengths: The paper is well organized, and the research is well motivated.
summary: The paper presents a hybrid architecture that combines Mamba-2 layers with softmax selective attention layers for long sequence modeling. The empirical study validates the effectiveness of the new model up to 1B model parameters.
weaknesses: 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.

Meta-review: The paper addresses an important challenge in efficient long-context language modeling with a hybrid architecture that combines state-space models and selective attention. While the proposed approach has merit, limited baseline comparisons, and insufficient empirical evidence reduce the overall impact, e.g. the experiments lack critical ablations of core components and analysis of memory and compute usage. I recommend rejecting the paper in its current form.

Decision: Reject