SSAM: Selective State Attention with Mamba

English | 中文

SSAM is a hybrid neural architecture The SSAM class in this repository (src/model/sudenmind.py) is a concrete language model instantiation of this architecture, named SudenMind.

---

Abstract

The central tension in long-sequence modeling is that Transformers rely on global attention with quadratic complexity, making efficient scaling difficult, while linear-complexity State Space Models (SSMs) excel at long-range modeling yet suffer from insufficient inductive bias for cross-segment global associations. This paper proposes Selective State Attention with Mamba (SSAM), a hybrid architecture that deeply integrates selective state space models with lightweight sparse attention. SSAM retains the linear-complexity advantage of parallel selective scanning inside each chunk, and performs sparse cross-chunk attention on chunk-terminal states to inject long-range context into the next layer via state initialization. This design reduces the computational cost of attention from quadratic in sequence length to quadratic in the number of chunks, while avoiding direct interference of traditional attention with token outputs—rendering attention a “navigation signal” for state transitions.

---

1. Introduction

Modern language models have evolved along two main paths:

1. Transformers achieve global dependencies between arbitrary positions through self-attention, at the cost of $O(T^2)$ computational and memory complexity.
2. State Space Models (SSMs), such as Mamba, compress sequence processing into $O(T)$ linear recurrences via selective scanning, greatly improving long-sequence efficiency. However, their cross-segment information transfer relies entirely on layer-wise compression of hidden states, potentially losing fine-grained long-range associations.

SSAM asks: Can we introduce cross-segment global attention at minimal cost while preserving the linear complexity of SSMs?

Our solution is to split the sequence into fixed-length chunks, apply parallel selective state space scanning inside each chunk, compute sparse attention on the terminal hidden state of each chunk, and feed the attention-weighted states as the initial state $h_{\text{init}}$ of the next SSM layer. Thus, attention no longer directly determines token outputs; instead, it modulates the state evolution of the next layer, balancing efficiency and global association.

---

2. Method

2.1 Overall Architecture

A decoder model based on the SSAM architecture (e.g., SudenMind) is stacked as follows:

Input tokens
    → Token Embedding
    → N × SSAMBlock
        (RMSNorm → SSAMCore → residual → optional FFN)
    → Final LayerNorm
    → LM Head → Logits

Each SSAMBlock consists of a pre-normalization residual SSAMCore and an optional feed-forward network (FFN). SSAMCore is the core computational unit of the SSAM architecture and adopts a dual-path design:

• Token output path: performs Mamba-style parallel selective scanning to produce per-token representations.
• State attention path: computes sparse cross-chunk attention on the terminal state of each chunk, producing aggregated state representations that serve as the initial state for the next layer.

The last layer skips the state attention path, since no subsequent layer needs to receive its output state.

2.2 Selective Scanning Inside Chunks

Given an input sequence $\mathbf{x} \in \mathbb{R}^{B \times T \times D}$, SSAMCore first splits it into two branches via a linear projection:

$$ \mathbf{x}_{\text{raw}}, \mathbf{z} = \text{split}(\mathbf{W}_{\text{in}} \mathbf{x}) $$

A causal depthwise separable 1D convolution followed by SiLU activation is then applied to $\mathbf{x}_{\text{raw}}$:

$$ \mathbf{x}_{\text{conv}} = \text{SiLU}(\text{CausalConv1d}(\mathbf{x}_{\text{raw}})) $$

The selective scan parameters $\Delta t$, $\mathbf{B}$, and $\mathbf{C}$ are projected, and $\Delta t$ is constrained by softplus:

$$ \Delta t = \text{softplus}(\mathbf{W}_{\Delta} \mathbf{x}_{\text{conv}}), \quad \mathbf{A} = -\exp(\mathbf{A}_{\log}) $$

The sequence is divided into $K = \lceil T / C \rceil$ chunks of length $C$, each executing a parallel selective scan:

$$ \mathbf{h}_t = \bar{\mathbf{A}}_t \mathbf{h}_{t-1} + \bar{\mathbf{B}}_t \mathbf{x}_t, \quad \mathbf{y}_t = \mathbf{C}_t^\top \mathbf{h}_t + \mathbf{D} \odot \mathbf{x}_t $$

The initial state $\mathbf{h}_{\text{init}}$ comes from the output of the previous layer’s state attention path, enabling inter-layer state transfer.

2.3 State-Level Sparse Cross-Chunk Attention

For each chunk $k$, the selective scan yields a terminal hidden state $\mathbf{s}k \in \mathbb{R}^{d{\text{inner}} \times d_{\text{state}}}$. To evaluate inter-chunk relationships at low cost, SSAM flattens $\mathbf{s}_k$ and projects it down to a compact vector $\mathbf{c}_k \in \mathbb{R}^{S}$ through a small MLP:

$$ \mathbf{c}_k = \text{MLP}_{\text{down}}(\text{flatten}(\mathbf{s}_k)) $$

Multi-head sparse attention is computed on these compact representations. To avoid quadratic complexity, attention is restricted to the $K$ chunks, yielding complexity $O(K^2 \cdot S) = O((T/C)^2 \cdot S)$, far below the $O(T^2 \cdot D)$ cost of token-level attention. Attention scores pass through a causal mask and a ReLU/Sparsemax-style normalization, producing an inter-chunk association distribution $\boldsymbol{\alpha} \in \mathbb{R}^{B \times H \times K \times K}$.

The key design is: attention is used only for scoring, while the weighted objects remain the original high-dimensional states. Specifically, for each target chunk $k$:

$$ \mathbf{w}_k = \sum_{j=1}^{k} \alpha_{kj} \cdot \text{flatten}(\mathbf{s}_j) $$

This avoids the information bottleneck caused by down-projection followed by up-projection, allowing the attention signal to act directly on the full state representation.

2.4 Gated State Fusion

To adaptively control the blending ratio between local states and globally aggregated states, SSAM introduces a scalar gate:

$$ g_k = \sigma(\mathbf{W}_g \cdot \text{flatten}(\mathbf{s}_k)) $$

The fused state is:

$$ \mathbf{f}_k = g_k \cdot \mathbf{w}_k + (1 - g_k) \cdot \text{flatten}(\mathbf{s}_k) $$

The gate is initialized near zero, so the model naturally degenerates to standard Mamba behavior at the beginning of training; as training progresses, the model gradually learns to exploit cross-chunk attention. The gate is also regularized to avoid extreme values (near 0 or 1) that would cause hard switching.

2.5 Independence of Token Output

An important design choice in SSAM is that the output of the state attention path does not participate in the current layer’s token output. The current layer’s token output is produced solely by the local selective scan:

$$ \mathbf{y} = \mathbf{W}_{\text{out}}(\mathbf{y}_{\text{ssm}} \odot \text{SiLU}(\mathbf{z})) $$

The state attention path output $\mathbf{f}k$ serves only as the initial state $\mathbf{h}{\text{init}}$ of the next SSM layer. This decoupling ensures that:

• token outputs retain local, efficient linear-recurrence properties;
• global attention acts only as a “navigation signal” for state evolution, without undermining the parallel training and linear inference advantages of SSMs.

---

3. Design Principles

3.1 State-Level Hybridization

SSAM’s hybridization occurs at the state level, not the token level. Transformers perform global attention on token representations; Mamba performs linear compression on hidden states; SSAM performs sparse selection on chunk-terminal states. This hierarchical hybridization preserves the linear scalability of SSMs while remedying their weak cross-segment association capability.

3.2 Low-Dimensional Scoring, High-Dimensional Weighting

Attention scoring is performed in a low-dimensional compact space to reduce cost, but the weighted objects remain the original high-dimensional states. This strategy avoids information loss caused by down/up projection while ensuring the scalability of the attention mechanism.

3.3 Soft Gating for Graceful Degradation

The gating mechanism allows the model to smoothly interpolate between “pure local SSM” and “global state attention.” At initialization the gate is near zero, degenerating the model to standard Mamba; during training the gate learns the optimal blending ratio.

3.4 Decoupled Attention from Token Prediction

Attention does not directly generate token outputs; it modulates the state initialization of the next layer. This decoupling allows SSAM to retain SSM linear complexity at inference time, adding only chunk-level sparse attention overhead.

---

4. Complexity Analysis

Component	Time Complexity	Space Complexity	Notes
Selective scan	$O(B \cdot T \cdot d_{\text{inner}} \cdot d_{\text{state}})$	$O(B \cdot K \cdot d_{\text{inner}} \cdot d_{\text{state}})$	Linear in sequence length
Sparse cross-chunk attention	$O(B \cdot K^2 \cdot S \cdot H)$	$O(B \cdot K^2 \cdot H)$	$K = T / C$, $S \ll d_{\text{inner}} d_{\text{state}}$
Down-projection MLP	$O(B \cdot K \cdot d_{\text{inner}} d_{\text{state}} \cdot S)$	$O(S \cdot d_{\text{inner}} d_{\text{state}})$	Parameter-level overhead

When $C$ is sufficiently large (e.g., $C = 128$), $K^2 = (T/C)^2$ is much smaller than $T^2$, strictly limiting the global attention overhead of SSAM. For very long sequences, the total complexity of SSAM approaches linear.

---

5. Implementation Notes

The SSAM architecture is implemented mainly in src/model/ssam.py. Core modules include:

• ChunkSparseAttention: chunk-level sparse attention with built-in RoPE positional encoding, $k$-value warm-up, and ReLU/Sparsemax-style normalization.
• SSAMCore: the core computational unit of the SSAM architecture, comprising selective scanning, state down-projection, sparse attention, and gated fusion.
• SSAMBlock: pre-normalization residual block with optional FFN.
• SSAMBackbone: decoder backbone.

The selective scan is implemented via a custom Triton kernel (src/model/triton_selective_scan.py) for efficient forward and backward passes, supporting initial-state injection. Normalization uses RMSNorm (src/model/norms.py).

The concrete model SudenMind based on the SSAM architecture is located in src/model/sudenmind.py; the full language model additionally includes token embedding, final LayerNorm, and the LM head.

---

6. Conclusion

SSAM proposes a hybrid architecture that fuses selective state space models with sparse attention at the state level. By restricting attention to chunk-terminal states and passing it to the next layer as an initial state, SSAM introduces cross-segment global association while preserving linear complexity. Its core contribution is redefining the role of attention in SSM architectures: attention is no longer a direct source of token outputs, but a navigation signal for state transitions. This design offers a new architectural option for building efficient long-sequence language models.

---

Reference Implementation

This repository provides a PyTorch implementation of the SSAM architecture, including a complete language model named SudenMind based on the architecture, covering pre-training, SFT fine-tuning, local chat, and OpenAI-compatible API deployment. For specific training and deployment instructions, please refer to the entry scripts and configuration files in the repository.