Hardware-aligned Hierarchical Sparse Attention for Efficient Long-term Memory Access

We sincerely appreciate your insightful feedback.

W1. Regarding Motivation, the authors do not design any sparse mechanism specifically for the Mamba block itself

We believe that "Enabling Random Long-Context Access for Mamba" is a goal, and the sparse mechanism for Mamba block could be one way to achieve it, but maybe not the only way. This paper, "via Hierarchical Sparse Attention," presents another approach to achieving this goal through a hybrid architecture.

As mentioned in our abstract and introduction, there is an information bottleneck in RNNs due to compressing variable-length sequences into fixed-dimensional representations, which inevitably leads to information loss. In contrast, the kv cache in attention grows with sequence length, which has no information bottleneck.

Thus, the current mainstream approach to enhancing Mamba's long-range random access ability is integrating Mamba with full attention [1]. As stated in lines 32-33 of the introduction, this leads to a trilemma. Consequently, our motivation lies in replacing full attention with sparse attention, as reflected in "Enabling Random Long-Context Access for Mamba via HSA."

Regarding Random Access

"Random access" refers to the ability to access any arbitrary past token, not just selecting chunks. For example, in an open-book exam, you locate a page (chunk) and then pinpoint specific details (tokens). Similarly, HSA enables access to any past token by first identifying the chunk and then the token-level details.

W4. What mechanism enables RAMba to achieve better extrapolation performance than both Mamba-2 and NSA?

We believe this is an interesting and important topic, so let us discuss it in advance.

First of all, the underperformance of Mamba-2 has been discussed in many studies [5]. The main reason is the information bottleneck issue, as we mentioned above.

To know why HSA achieves better extrapolation performance, let's take a close look at how NAS and HSA work in the following example. Assuming every 2 tokens form a chunk, separated by |:

chunk ids             |       0      |      1    |      2    |     3     | ...
token positions       |   0   |  1   |  2  |  3  |  4  |  5  |  6  |  7  | ... , i <- the currrent token
attn logits[i, *]     | -20.0 | 20.0 | 1.0 | 0.0 | 2.0 | 1.0 | 1.0 | 1.0 |
attn weights (softmax)|   0.0 |  1.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |  // O(L$^2$), L: sequence length
chunk attn logits sum |      0.0     |    0.5    |    1.5    |    1.0    |  // O(L*(L/S), S: chunk size

NSA selects chunks according to "chunk attn logits sum" for efficiency. If only the top-2 chunks can be selected, NSA would choose chunks 2 and 3, thus missing the chunk 0 with the highest sum of attention weights, leading to inaccurate chunk selection.

Assuming NSA selects the top-2 chunks for sparse attention, the following process would be:

chunk ids         |      2      |      3      | 
token positions   |   4  |   5  |   6  |   7  |
attn logits       |  2.0 |  1.0 |  1.0 |  1.0 |
                  |<-----     softmax   ----->|
attn weights      | 0.49 | 0.17 | 0.17 | 0.17 |

Throughout this process, chunk selection can not be explicitly learned.

For the same example, HSA works as follows:

chunk ids          |       0      |      1     |      2    |     3     | ...
token positions    |   0   |   1  |  2   |  3  |  4  |  5  |  6  |  7  | ...
attn logits        | -20.0 | 20.0 |  1.0 | 0.0 | 1.0 | 1.0 | 1.0 | 1.0 | ...
s[i,c]             |    s [i,0]   |  s [i,1]   |  s [i,2]  |  s [i,3]  | ...

Where s [i,c] represents learnable token-to-chunk relevance scores as described in Line 155 of the paper. For example, HSA may also select chunks 2 and 3 when parameters are randomly initialized, we have:

chunk ids            |           2           |          3          |
token positions      |     4     |     5     |     6    |     7    |
attn logits          |   2.0     |    1.0    |    1.0   |    1.0   |
intra-chunk weighting|<--     softmax      ->|<--    softmax     ->|
token-level weights  |   0.73    |    0.27   |    0.5   |    0.5   |
s[i,c]               |        s[i,2]         |        s[i,3]       |
inter-chunk weighting|<---            stick-breaking           --->|
chunk-level weights  |        w[i,2]         |        w[i,3]       |   // s.t. w [i,2} + w [i,3] <= 1
final weights        |0.73×w[i,2]|0.27×w[i,2]|0.5×w[i,3]|0.5×w[i,3]|

In this process, w[i,c], derived from s[i,c], participates in the final attention weight computation and the whole forward pass. This allows it to receive gradients and allocate higher weights to chunks with more important tokens. Even if random initialization initially misses the most relevant chunk 0, continuous training enables the model to gradually learn to select the most relevant chunks.

Such a learnable chunk selection mechanism is essentially a kind of token-to-chunk retrieval, which may lead to better extrapolation performance

W2. Is there any literature supporting the claim that augmenting RNNs with attention mechanisms leads to poor length extrapolation?

Yes. Specifically, lines 29-30 refer to hybrid architectures like Mamba with full attention. Mamba with sliding window attention is not included in this discussion. The issue stems from the poor length extrapolation ability of full-attention, which has never been fundamentally solved.

For instance, Figure 1 of YaRN [4] shows how perplexity (ppl) increases when full-attention-based models with RoPE extrapolate to unseen context lengths.

Table 2 in [7] also reveals that, even with temperature scaling and discard positional encodings, perplexity still dramatically increases beyond a certain length, indicating poor extrapolation.

The existing full-attention-based models with 1 million context lengths primarily rely on techniques like ring-attention to brute-force scale the pretraining length, which requires huge computational resources.

As minimax[6] mentioned in their abstract, "The context window of MiniMax-Text-01 can reach up to 1 million tokens during training and extrapolate to 4 million tokens during inference at an affordable cost. ” They also acknowledge the limitation of only being able to extrapolate up to 4 times the training length. The 1 million pre-training context window is also achieved via ring-attention.

Our work achieves extrapolation from 4K to 64M (16000×) in passkey retrieval and maintains high performance in both training and inference.

W3. Does the introduction of the transformer-based bi-directional encoder and the mean-pooled representation compromise the model's causality?

No, it doesn't. Each token cannot access the chunk it belongs to, it can only access the previous chunks. We will make further clarification in the revised version.

W5. Can HSA operate independently of the Mamba blocks? If a model were designed using only HSA, would it still be faster than NSA?

We believe so. It can work together with sliding window attention. We believe HSA would still be faster than NSA because we do not require token compression attention, thus the overall FLOPs are lower.

Q1. Why does HSA adopt the stick-breaking attention mechanism for constructing chunk weights, rather than other forms of attention with positional encoding? I would like to see a theoretical justification or supporting ablation studies.

If we use RoPE, a straightforward issue would be the out-of-domain problem during extrapolation (when the input length exceeds the pretraining length), as discussed in [4] and [7].

Here is an ablation study to replace the stick-breaking mechanism with softmax+RoPE:

ppl with context length=64K (the lower the better):

RULER: (context len=64K) (the higher the better)

Q2. Does the equation for the HSA layer in this line omit the Mamba layers?

No, sorry for the confusion. It doesn't omit the Mamba layers. The layer indices here include both Mamba and HSA layers, and we will clarify this point further later.

Q3. Why does NSA's implementation currently not support generation?

As there is no official implementation for NSA thus we use an implementation from the open-source community (line 230).

NSA's batch forward mode (used for training) calculates next-token perplexity for full sequences without needing KV cache management for new tokens, as no new tokens are generated. In contrast, generation mode requires handling a variable number of newly generated tokens, complex kv cache management, and positional encoding, especially with different attention mechanisms. This complexity is why vLLM performs well for inference but not training, while Megatron is optimized for the reverse. The KV management for generation was incomplete at the time of the paper's submission.

Even so, this does not undermine our conclusions. The majority of experiments can be evaluated via batch forward passes, and the performance on PPL and RULER is sufficient to support our claims.

[1] Jamba: Hybrid Transformer-Mamba Language Models

[2] SAMBA: SIMPLE HYBRID STATE SPACE MODELS FOR EFFICIENT UNLIMITED CONTEXT LANGUAGE MODELING

[3] Understanding and Improving Length Generalization in Recurrent Models

[4] YaRN: Efficient Context Window Extension of Large Language Models

[5] Repeat After Me: Transformers are Better than State Space Models at Copying Transformers are Better than State Space Models at Copying

[6] MiniMax-01: Scaling Foundation Models with Lightning Attention

[7] Length Generalization of Causal Transformers without Position Encoding

We hope our replies address your concerns, and we would appreciate it very much if you could give us more support.

Thank you very much for your valuable suggestions and for recognizing the contributions of HSA.

We greatly appreciate your feedback, particularly the suggestion to move the more technical content into the appendix, which we find highly helpful. Given the positive feedback on the example, we plan to present it as a figure within the methodology section to better illustrate the distinction between NSA and HSA, replacing the original kernel design section. The kernel design section will then be relocated to the appendix. This adjustment not only enhances the overall flow but also provides clearer reasoning behind issues such as "inaccurate chunk selection."

The space freed up by this reorganization will be utilized to address the ambiguous phrases you highlighted and to provide a more comprehensive explanation of the design rationales behind each component, with particular elaboration on the encoder.

We sincerely appreciate your insightful feedback and are delighted to hear that you largely enjoyed reading our paper.

W1, Q1: Can you provide concrete evidence supporting the claim that existing chunk-based methods "suffer from inaccurate chunk selection"?

Apologies for the confusion.

In line 38, "a closer examination reveals a critical weakness" refers to the evidence observed in our experiments. The chunk-based methods mentioned in our introduction mainly refer to NSA, which we consistently compare against in all experiments. Notably, in Table 3, Mamba(w/ NSA) performs significantly worse than RAMba(w/o enc) on the more challenging MQ-NIAH task, for in-domain sequence lengths. The gap becomes even more pronounced in length extrapolation. Please note these two models only differ in the attention module (NSA vs HSA).

The results of RAMba(w/ NSA, w/ m.r.) further support the claim. Without RNN's memory, where retrieval relies solely on NSA, there is a significant performance drop across all tasks. In contrast, RAMba(w/ m.r.) exhibits no such issue.

Here, please allow us to further explain through theory and specific examples why NSA may suffer from inaccurate chunk selection, as well as how HSA helps mitigate this problem.

The core principle of NSA's learnable chunk selection lies in approximating the token-to-chunk relevance using unnormalized attention scores. In self-attention, the relevance of token j to token i is defined as:

p[i,j] = $\frac{e^{logits_{i,j}}}{Z_i}, \quad Z_i = \sum_{j < i}{e^{logits_{i,j}}}$,

where logits[i,j] = $q_i^T k_j$ is the dot product between the query $q_i$ of token i and the key $k_j$ of token j.

The relevance of chunk c to token i is ideally the sum of the relevance of all tokens within that chunk:

r[i,c] = $\sum_{j \in \mathcal{C}}{p_{i,j}}= \frac{1}{Z_i} \sum_{j \in \mathcal{C}} e^{logits_{i,j}}$,

where $\mathcal{C}$ denotes tokens in chunk c. However, calculating r[i,c] requires computing $Z_i$ (the full softmax normalization across all tokens), which would necessitate full $Q^TK$ computation and thus undermine the computational efficiency.

To ensure efficiency, NSA approximates chunk relevance using the mean-pooled key representation $K_c$ of chunk c:

r'[i,c] = $q_i^T K_c= q_i^T \frac{1}{S} \sum_{j \in \mathcal{C}}{k_j} = \frac{1}{S} \sum_{j \in \mathcal{C}} q_i^T k_j = \frac{1}{S} \sum_{j \in \mathcal{C}}{logits_{i,j}}$,

where $K_c=\frac{1}{S} \sum_{j \in \mathcal{C}}{k_j}$. While this approximation bypasses the need for normalized softmax scores across all tokens, it introduces a discrepancy between r'[i,c] (the approximation) and r[i,c] (the ideal one).

Let's consider the following example, assuming every 2 tokens form a chunk, separated by |:

chunk ids       |       0      |      1    |      2    |     3     | ...
token positions |   0   |  1   |  2  |  3  |  4  |  5  |  6  |  7  | ... , i
logits[i,j]     | -20.0 | 20.0 | 1.0 | 0.0 | 2.0 | 1.0 | 1.0 | 1.0 |
p[i,j]          |   0.0 |  1.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |
r[i,c]          |      1.0     |    0.0    |    0.0    |    0.0    |
r'[i,c]         |      0.0     |    0.5    |    1.5    |    1.0    |

If only the top-2 chunks can be selected, NSA would choose chunks 2 and 3 according to r' [i,c], thus missing the chunk 0 with the highest sum of attention weights.

Assuming NSA selects the top-2 chunks for sparse attention, the following process would be:

chunk ids         |      2      |      3      | 
token positions   |   4  |   5  |   6  |   7  |
logits[i,j]       |  2.0 |  1.0 |  1.0 |  1.0 |
                  |<-----     softmax   ----->|
p'[i,j]           | 0.49 | 0.17 | 0.17 | 0.17 |

Throughout this process, r'[i,c] only participates in chunk selection but not in the forward computation, nor does it receive gradients.

The NSA's inaccurate chunk selection issue stems from using unnormalized attention logits to estimate chunk importance, with r'[i,c] not learnable, making the inaccuracy of chunk selection unavoidable.

For the same example, HSA works as follows:

chunk ids          |       0      |      1     |      2    |     3     | ...
token positions    |   0   |   1  |  2   |  3  |  4  |  5  |  6  |  7  | ...
logits[i,j]        | -20.0 | 20.0 |  1.0 | 0.0 | 1.0 | 1.0 | 1.0 | 1.0 | ...
p[i,j]             |  0.0  |  1.0 |  0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | ...
s[i,c]             |    s [i,0]   |  s [i,1]   |  s [i,2]  |  s [i,3]  | ...

Where s [i,c] represents learnable token-to-chunk relevance scores as described in Line 155 of the paper. For example, HSA may also select chunks 2 and 3 when parameters are randomly initialized, we have:

chunk ids            |           2           |          3          |
token positions      |     4     |     5     |     6    |     7    |
attn_logits [i,j]    |   2.0     |    1.0    |    1.0   |    1.0   |
intra-chunk weighting|<--     softmax      ->|<--    softmax     ->|
token-level weights  |   0.73    |    0.27   |    0.5   |    0.5   |
s[i,c]               |        s[i,2]         |        s[i,3]       |
inter-chunk weighting|<---            stick-breaking           --->|
chunk-level weights  |        w[i,2]         |        w[i,3]       |   // s.t. w [i,2} + w [i,3] <= 1
final weights        |0.73×w[i,2]|0.27×w[i,2]|0.5×w[i,3]|0.5×w[i,3]|

In this process, w[i,c], derived from s[i,c], participates in the final attention weight computation and the whole forward pass. This allows it to receive gradients and allocate higher weights to chunks with more important tokens. Even if random initialization initially misses the most relevant chunk 0, continuous training enables the model to gradually learn to select the most relevant chunks.

This point can be further validated through a pure NSA Transformer. RULER results:

W2. The multiple interacting components make it difficult to isolate individual contributions

We fully agree with your perspective and have conducted ablation studies on each component to observe their individual impacts.

Table 3 includes all ablation experiments, such as removing the encoder (w/o enc), memory reset (w/o m.r.), and replacing HSA with NSA or purely MAMBA, reflecting the influence of different modules on the overall system.

The most straightforward comparison is between RAMba(w/o enc), Mamba, and Mamba(w/ NSA). The only difference lies in HSA, which shows significant improvement for inputs of 64K and below. The other variants, such as memory reset and encoder, primarily serve to further enhance HSA's extrapolation capability.

W3. NSA's implementation currently does not support generation.

As there is no official implementation for NSA thus we use an implementation from the open-source community (line 230).

NSA's batch forward mode (used for training) calculates next-token perplexity for full sequences without needing KV cache management for new tokens, as no new tokens are generated. In contrast, generation mode requires handling a variable number of newly generated tokens, complex kv cache management, and positional encoding, especially with different attention mechanisms. This complexity is why vLLM performs well for inference but not training, while Megatron is optimized for the reverse. The KV management for generation was incomplete at the time of the paper's submission.

Even so, this does not undermine our conclusions. The majority of experiments can be evaluated via batch forward passes, and the performance on PPL and RULER is sufficient to support our claims.

W4. I find it hard to accept the claim that Mamba "excels".

Thanks for pointing that out. The original phrasing was imprecise. We will revise it to: "Mamba performs better in long sequential statistical tasks."

W6/Q5. How memory offloading can provide empirical benefits. Regarding computational complexity analysis.

Sorry for the confusion, the "memory offloading" refers to offloading the KV cache from GPU to the CPU, not HBM to SRAM; Memory offloading does not provide computational benefits; in fact, it slows down generation but reduces the GPU's memory footprint, as reflected in Figure 4(3) and Table 7. During long-sequence inference, the KV cache could overwhelm GPU memory. Offloading the KV cache to CPU and just reloading the KV cache of selected chunks can largely save the memory footprint. Therefore, memory offloading has no relationship with whether HSA is memory-bound or compute-bound.

The complexity of HSA is as follows:

L: sequence length, d: hidden size, K: number of selected chunks, S: number of tokens in a chunk Retrieval: $O(L \times \frac{L}{S} \times d)$

Attention: $O(L \times K \times S \times d)$ Compared to your version, there is no $S^2$, each token attends to $K \times S$ tokens.

Chunk-level fusion: $O(L \times K \times d)$

Q2: Do the PPUs support tensor cores?

Yes, they do. Here, the "PPU" is just another name for GPUs not produced by NVIDIA. It also supports Torch and Triton. Recently, we have access to the H800, and the throughput of 8*H800 is twice the results shown in Table 6, which aligns with Figure 4. Without tensor cores, the PPU could not achieve such throughput.

Q3. Are the perplexity results on validation or test sets?

On validation sets.

Q4. How and where is the sorting operation applied?

When computing token-to-chunk weighting in line 155. It's outside the HSA triton kernel; weights and indices are already computed before conducting HSA.

Your suggestion on discussing the chunk selection issue is very helpful, and we will include it in the Appendix. We hope our replies address your concerns, and we would appreciate it very much if you could consider raising the score.

We sincerely appreciate your insightful feedback and recognition of our work.

W1,Q1. Regarding the model’s architectural complexity and its new hyperparameters (S, K, G): could the authors comment on their sensitivity? How were the reported values chosen, and what is the expected impact of varying them?

Based on previous research on hybrid models ([1], Fig.4), the optimal ratio between attention and Mamba is generally around 8.25%~20%. Therefore, we inserted 3 HSA layers into the 48-layer Mamba model. Generally, S is set to 64 to better align with hardware.

The value of K primarily impacts computational cost—larger K values typically result in better performance but also higher computational load and lower sparsity. The optimal value of K is worth conducting an empirical study on. However, since this work requires comparison with other baselines to verify its effectiveness, we have not delved into this aspect here. This will be left for future work.

W2,Q2. Regarding the memory reset mechanism: could the authors provide an ablation study against simpler strategies (e.g., resetting to zero or to the last state) to better justify the choice of a randomized approach?

Thank you for your insightful question regarding the memory reset mechanism. To address your concern, we would like to clarify that our approach involves initializing the memory state using the last hidden state from a randomly selected segment.

In fact, a recent study [2] has conducted an empirical analysis specifically on this topic. They refer to using the last hidden state from a random segment as the initial state as "state passing," and using the last hidden state from the preceding sentence as the initial state as "truncated backpropagation through time (BPTT or TBTT). They compared different strategies, including setting to zero, TBTT, and state passing. Their results ([2], Fig. 5) demonstrate that initializing the initial state with the last state from a random segment (as employed in our paper) consistently outperforms other alternatives when extrapolating, such as resetting to zero or initializing with a randomized vector. This evidence supports the effectiveness of our chosen approach.

W3, Q3. We suggest the authors provide results for an NSA baseline augmented with the same chunk encoder

We sincerely appreciate your suggestion. We did not conduct the comparison for two reasons. First, we believe that RAMba (w/o enc) is a fair comparison with Mamba w/ NSA. Secondly, NSA is inherently incompatible with the chunk encoder, as it is unique to the hierarchical attention framework. Please allow us to elaborate further.

One purpose of introducing RAMba (w\o enc) is to enable a fair comparison with NSA. Aside from the differences in the attention module, RAMba (w\o enc) and Mamba (w\ NSA) are identical in every other aspect. However, RAMba (w\o enc) significantly outperforms Mamba (w\ NSA) on RULER (Table. 3), and this comparison validates that the performance improvement mainly stems from HSA.

In addition, NSA itself is not compatible with the encoder. Let's take an example to see the difference between how NSA and HSA work.

Assuming every 2 tokens form a chunk, separated by |:

chunk ids             |       0      |      1    |      2    |     3     | ...
token positions       |   0   |  1   |  2  |  3  |  4  |  5  |  6  |  7  | ... , i <- the currrent token
attn logits[i, *]     | -20.0 | 20.0 | 1.0 | 0.0 | 2.0 | 1.0 | 1.0 | 1.0 | 
attn weights (softmax)|   0.0 |  1.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |  // O(L$^2$), L: sequence length
chunk attn logits sum |      0.0     |    0.5    |    1.5    |    1.0    |  // O(L*(L/S), S: chunk size

NSA selects chunks according to "chunk attn logits sum" for efficiency. If only the top-2 chunks can be selected, NSA would choose chunks 2 and 3, thus missing the chunk 0 with the highest sum of attention weights, leading to inaccurate chunk selection.

Assuming NSA selects the top-2 chunks for sparse attention, the following process would be:

chunk ids         |      2      |      3      | 
token positions   |   4  |   5  |   6  |   7  |  // sentence-level RoPE applied here (0 ~ L)
attn logits       |  2.0 |  1.0 |  1.0 |  1.0 | 
                  |<-   enc?  ->|<-   enc?  ->|  
                  // Chunk-wise encoding here will disrupt global positional information.
                  |<-----     softmax   ----->| 
attn weights      | 0.49 | 0.17 | 0.17 | 0.17 |

If chunk encoding is applied at the chunk level, it will destroy the global positional information.

For the same example, HSA works as follows:

chunk ids          |       0      |      1     |      2    |     3     | ...
token positions    |   0   |   1  |  2   |  3  |  4  |  5  |  6  |  7  | ...
attn logits        | -20.0 | 20.0 |  1.0 | 0.0 | 1.0 | 1.0 | 1.0 | 1.0 | ...
s[i,c]             |    s [i,0]   |  s [i,1]   |  s [i,2]  |  s [i,3]  | ...

Where s [i,c] represents learnable token-to-chunk relevance scores as described in Line 155 of the paper. For example, HSA may also select chunks 2 and 3 when parameters are randomly initialized, we have:

chunk ids            |           2           |          3          |
                     |<---      enc      --->|<---      enc    --->| // chunk-wised encoding
attn logits          |   2.0     |    1.0    |    1.0   |    1.0   | // intra-chunk RoPE only (0~S for each chunk)
intra-chunk weighting|<--     softmax      ->|<--    softmax     ->|
token-level weights  |   0.73    |    0.27   |    0.5   |    0.5   |
s[i,c]               |        s[i,2]         |        s[i,3]       |
inter-chunk weighting|<---            stick-breaking           --->| // with positional inductive bias
chunk-level weights  |        w[i,2]         |        w[i,3]       |   // s.t. w [i,2} + w [i,3] <= 1
final weights        |0.73×w[i,2]|0.27×w[i,2]|0.5×w[i,3]|0.5×w[i,3]|

Due to the positional inductive bias inherent in the stick-breaking attention mechanism, it becomes feasible to employ a bidirectional encoder at the chunk level.

Moreover, in this process, w[i,c], derived from s[i,c], participates in the final attention weight computation and the whole forward pass. This allows it to receive gradients and allocate higher weights to chunks with more important tokens. Even if random initialization initially misses the most relevant chunk 0, continuous training enables the model to gradually learn to select the most relevant chunks.

[1] An Empirical Study of Mamba-based Language Models

[2] Understanding and Improving Length Generalization in Recurrent Models

We sincerely appreciate your insightful feedback.

Q3,W1. I did not figure out the difference between HSA and NSA. How "chunk-aware" learning is achieved beyond standard backpropagation through the hierarchical structure could be elaborated upon

This is the most important contribution of this work, so please allow us to explain it first.

Why NSA suffer from inaccurate chunk selection, and why can HSA learn from back-propagation?

The core principle of NSA's learnable chunk selection lies in approximating the token-to-chunk relevance using unnormalized attention scores. In self-attention, the relevance of token j to token i is defined as:

p[i,j] = $\frac{e^{logits_{i,j}}}{Z_i}, \quad Z_i = \sum_{j < i}{e^{logits_{i,j}}}$,

where logits[i,j] = $q_i^T k_j$ is the dot product between the query $q_i$ of token i and the key $k_j$ of token j.

The relevance of chunk c to token i is ideally the sum of the relevance of all tokens within that chunk:

r[i,c] = $\sum_{j \in \mathcal{C}}{p_{i,j}}= \frac{1}{Z_i} \sum_{j \in \mathcal{C}} e^{logits_{i,j}}$,

where $\mathcal{C}$ denotes tokens in chunk c. However, calculating r[i,c] requires computing $Z_i$ (the full softmax normalization across all tokens), which would necessitate full $Q^TK$ computation and thus undermine the computational efficiency.

To ensure efficiency, NSA instead approximates chunk relevance using the mean-pooled representation of keys:

r'[i,c] = $q_i^T K_c= q_i^T \frac{1}{S} \sum_{j \in \mathcal{C}}{k_j} = \frac{1}{S} \sum_{j \in \mathcal{C}} q_i^T k_j = \frac{1}{S} \sum_{j \in \mathcal{C}}{logits_{i,j}}$,

where $K_c$ represents the mean-pooling of key representations within chunk c. While this approximation bypasses the need for normalized softmax scores across all tokens, it introduces a discrepancy between r'[i,c] (the approximation) and r[i,c] (the ideal one).

Let's consider the following example, assuming every 2 tokens form a chunk, separated by |:

chunk ids       |       0      |      1    |      2    |     3     | ...
token positions |   0   |  1   |  2  |  3  |  4  |  5  |  6  |  7  | ... , i (current token)
logits[i,j]     | -20.0 | 20.0 | 1.0 | 0.0 | 2.0 | 1.0 | 1.0 | 1.0 |
p[i,j]          |   0.0 |  1.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |
r[i,c]          |      1.0     |    0.0    |    0.0    |    0.0    |
r'[i,c]         |      0.0     |    0.5    |    1.5    |    1.0    |

If only the top-2 chunks can be selected, NSA would choose chunks 2 and 3 according to r' [i,c], thus missing the chunk 0 with the highest sum of attention weights.

Assuming NSA selects the top-2 chunks for sparse attention, the following process would be:

chunk ids         |      2      |      3      | 
token positions   |   4  |   5  |   6  |   7  |
logits[i,j]       |  2.0 |  1.0 |  1.0 |  1.0 |
                  |<-----     softmax   ----->|
p'[i,j]           | 0.49 | 0.17 | 0.17 | 0.17 |

Throughout this process, r'[i,c] only participates in chunk selection but not in the forward computation, nor does it receive gradients.

The NSA's inaccurate chunk selection issue stems from using unnormalized attention logits to estimate chunk importance, with r'[i,c] not learnable, making the inaccuracy of chunk selection unavoidable.

For the same example, HSA works as follows:

chunk ids          |       0      |      1     |      2    |     3     | ...
token positions    |   0   |   1  |  2   |  3  |  4  |  5  |  6  |  7  | ...
logits[i,j]        | -20.0 | 20.0 |  1.0 | 0.0 | 1.0 | 1.0 | 1.0 | 1.0 | ...
p[i,j]             |  0.0  |  1.0 |  0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | ...
s[i,c]             |    s [i,0]   |  s [i,1]   |  s [i,2]  |  s [i,3]  | ...

Where s [i,c] represents learnable token-to-chunk relevance scores as described in Line 155 of the paper. For example, HSA may also select chunks 2 and 3 when parameters are randomly initialized, we have:

chunk ids            |           2           |          3          |
token positions      |     4     |     5     |     6    |     7    |
attn_logits [i,j]    |   2.0     |    1.0    |    1.0   |    1.0   |
intra-chunk weighting|<--     softmax      ->|<--    softmax     ->|
token-level weights  |   0.73    |    0.27   |    0.5   |    0.5   |
s[i,c]               |        s[i,2]         |        s[i,3]       |
inter-chunk weighting|<---            stick-breaking           --->|
chunk-level weights  |        w[i,2]         |        w[i,3]       |   // s.t. w [i,2} + w [i,3] <= 1
final weights        |0.73×w[i,2]|0.27×w[i,2]|0.5×w[i,3]|0.5×w[i,3]|

In this process, w[i,c], derived from s[i,c], participates in the final attention weight computation and the whole forward pass. This allows it to receive gradients and allocate higher weights to chunks with more important tokens. Even if random initialization initially misses the most relevant chunk 0, continuous training enables the model to gradually learn to select the most relevant chunks.

Q1. Can the authors also report the results for pure NSA w/o SWA as used in MiniCPM4?

Since MiniCPM4 had not been released at the time of submitting the paper, we did not include a comparison. However, this suggestion was very helpful. We tested it, and the results are as follows:

Perplexity:

RULER:

This result further validates our argument that the NSA has issues with chunk selection, even for in-domain length.

W1. Whether the primary performance gains stem from the stick-breaking mechanism itself, rather than HSA.

RAMba$_{w/o s.b}$ means using softmax without position encoding instead of the stick-breaking mechanism (please refer to Line 242). We believe that the performance degradation is due to the inability to distinguish the relative distance between chunks, rather than additional improvement brought by the stick-breaking mechanism.

Regarding the extrapolation ability, the experiments in the original stick-breaking attention paper[1] (Figure 6) show that stick-breaking attention itself cannot resolve the NIAH extrapolation issue, which demonstrates that the extrapolation capability is not brought by the stick-breaking mechanism. Moreover, the stick-breaking mechanism is not applied to tokens, but only to chunks, as we mentioned above.

In summary, stick-breaking is one component of the overall design, and it undoubtedly contributes to our results to some extent. However, the core contributions of this paper, the extrapolation capability and the high performance of sparse attention, are still primarily achieved by HSA.

W2. Justification for the Transformer-based Encoder.

Since the intra-chunk attention of HSA is conducted with each chunk separately, an intuitive idea is to use a bi-directional encoder to enhance chunk-level representations. Please note that it doesn't break causality, as each token is not allowed to access the chunk it belongs to.

Using bi-directional Mamba layers may be another choice. Here are the results for RULER:

The results are quite close, which validates that the choice of encoder has minimal impact.

Q2. Could you elaborate more on BPTT?

Existing parallelizable RNNs often suffer from Oversized States[3]. A common solution is to use BPTT (Backpropagation Through Time) for pretraining.

Simply put, the last state from the previous step is used, and after gradient truncation, it becomes the initial state for the current step.

This aligns with the recent work[4] by the author of Mamba, where they term this mechanism as 'state passing.'

As demonstrated in the experiments in [4] (Fig. 5), state passing has a negligible impact on results but significantly enhances extrapolation capability.

Q5. Phonebook lookup

Here are the results:

Q4. MK-NIAH on RULER

The MQ-NIAH test in Table 3 covers the multi-key NIAH test. It contains 6 key-value pairs in the context. Moreover, the phonebook lookup task is a harder version of MK-NIAH.

[1] Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study

[2] Retrieval-Pretrained Transformer: Long-range Language Modeling with Self-retrieval

[3] Stuffed Mamba: Oversized States Lead to the Inability to Forget

[4] Understanding and Improving Length Generalization in Recurrent Models

Your suggestion to compare Transformers with the NSA is extremely helpful. We have provided the relevant data to address your concerns and hope our responses adequately resolve the issues. We would greatly appreciate it if you could reconsider and raise the score.

Thank you for your very thoughtful responses, which have addressed most of my concerns. I also greatly appreciate the authors' hard work in recent days. However, I still have some reservations regarding the complexity of the design and how it scales. As such, I am maintaining my score of 3. I am, however, very much looking forward to the improved version of the HSA.

Thank you very much for your prompt reply.

Regarding scalability:

1. Tables 6-7 & Figure 4 comprehensively demonstrate our model's scalability in terms of context length (16K~64K) and parameter size (370M~3B).
2. Table 5 further reports results after continued pretraining on a 3B model.

If your concerns are about latency and throughput for short sequences, here are the results:

Per-step inference cost:

Even for short texts, the additional overhead introduced by HSA is less than 10% compared to vanilla Mamba.

Training throughputs at 2K context length:

For short context lengths, the throughput of HSA is also comparable to both Mamba and Mamba with full attention. Fundamentally, the design of HSA does not introduce any drawbacks for short-text modeling.

In summary, our experimental results cover context lengths from 64 to 32K and model parameters from 370M to 3B, and there is no objective evidence indicating any significant scalability deficiencies in HSA.

Regarding design complexity, compared to the mainstream Mamba+attention approach, we replace full-attention with hierarchical sparse attention (HSA), with the only additional design being the chunk-wise encoder, whose necessity we have analyzed in the ablation study. Additionally, memory reset is not part of the model architecture but merely a training technique. Even without the memory reset, our model still outperforms other baselines. We aim to share more empirical insights with the community, even though it introduces some understanding overhead.

As a scientific study, we strive to analyze issues based on objective data. We have provided all the supporting data for our claims and analyzed the necessity of each design component. If there are additional objective results that could address your remaining concerns, we welcome further discussion.

Dear Reviewer fVSb

As suggested by Reviewer XZCu, we consider moving the kernel design section to the appendix, freeing up space to provide more detailed explanations of our existing design rationales.

Why the chunk-wise Encoder is Necessary:
Each chunk requires a representation to be retrieved, which can be achieved through either a parameter-free or a parameterized approach. The parameter-free approach is mean pooling of hidden states or keys inside a chunk, while the parameterized approach introduces an additional encoder. We empirically find that using an extra encoder can significantly improve extrapolation performance. Thus, we included such a design in the paper along with an ablation study to validate its effectiveness.

Why Alternate One HSA Layer with G Mamba Layers in the Upper Decoder:
Current hybrid architectures combining Mamba and attention layers [1] typically alternate X layers of Mamba with one layer of attention. [1] indicates there's an optimal ratio exists, and excessive attention layers may lead to performance degradation under the same parameter budget. Therefore, our choice to alternate attention with multiple Mamba based on their empirical results.

The modules we employ (Mamba, HSA, and chunk-wise encoder) are all essential. The proposed model architecture is a deliberate design supported by prior findings and our empirical results. We hope this clarification addresses concerns regarding complexity and provides a clearer understanding of the rationale behind our architecture.

[1] An Empirical Study of Mamba-based Language Models