Introduction
The quest for efficient sequence modeling has led to significant innovations beyond Transformers. While Mamba introduced Selective State Space Models (SSSM) as a promising alternative to attention mechanisms, its scalability across diverse applications remained limited. S-Mamba addresses this challenge by introducing a suite of scalable selective state space models that achieve superior performance across language modeling, time series forecasting, vision tasks, and more.
This article explores the S-Mamba architecture, its innovations, and its applications in modern AI systems.
The Evolution from Mamba to S-Mamba
Mamba’s Core Innovation
To understand S-Mamba you first need to understand what Mamba changed. Sequence models based on attention have an inherent scaling problem: every token in a sequence attends to every other token, so memory and computation grow quadratically with sequence length. State space models (SSMs) sidestep this by processing tokens through a fixed-size recurrent state, giving linear-time inference, but early SSMs used input-independent parameters that limited their ability to selectively attend to relevant information. Mamba’s core contribution was making the state space parameters depend on the input itself, which the paper calls selectivity, so the model can choose what to remember and what to forget at each step.
The block below captures the essential machinery that makes this work. The input is projected and split into two streams: one that flows through the state space computation and a second, gating stream that modulates the result. A lightweight convolution provides local context before the selective parameters are computed, and the output is gated with the silu-activated second stream before being projected back to the model dimension. Every parameter here that is marked as input-dependent, the projection matrices that produce B, C, and dt, is where the selectivity lives, and that is the difference between a generic SSM and a Mamba-style one.
class MambaBlock:
"""
Original Mamba: Selective State Space Model
"""
def __init__(self, d_model, d_state=128, d_conv=4, expand=2):
self.d_model = d_model
self.d_state = d_state
self.d_conv = d_conv
self.d_inner = int(expand * d_model)
# Input projection
self.in_proj = nn.Linear(d_model, self.d_inner * 2)
# Convolutional layer for local context
self.conv1d = nn.Conv1d(
self.d_inner,
self.d_inner,
kernel_size=d_conv,
padding=d_conv - 1,
groups=self.d_inner
)
# SSM parameters (selective)
self.x_proj = nn.Linear(self.d_inner, d_state * 2)
self.dt_proj = nn.Linear(self.d_inner, self.d_inner)
# State space parameters
self.A_log = nn.Parameter(torch.randn(self.d_inner, d_state))
self.D = nn.Parameter(torch.ones(self.d_inner))
# Output projection
self.out_proj = nn.Linear(self.d_inner, d_model)
def forward(self, x):
"""
Mamba forward pass with selective mechanism
"""
# Project input
xz = self.in_proj(x)
x_inner, z = xz.chunk(2, dim=-1)
# Convolution for local context
x_conv = self.conv1d(x_inner.transpose(1, 2))
x_conv = x_conv[:, :, :-self.d_conv + 1].transpose(1, 2)
# Selective SSM: compute parameters based on input
ssm_params = self.x_proj(x_conv)
B, C = ssm_params.chunk(2, dim=-1)
# Discretize continuous parameters
dt = F.softplus(self.dt_proj(x_conv))
# State space computation (selective)
# This is where Mamba differs: parameters depend on input
y = self.selective_state_space(x_conv, dt, A, B, C, D)
# Gating mechanism
y = y * F.silu(z)
# Output projection
return self.out_proj(y)
A few implementation details in this forward pass are worth studying because they quietly enable the
whole architecture. The convolution is depthwise, meaning each input channel is convolved
independently, which adds local context without the parameter blow-up of a full convolution. The
A_log parameter is stored in log space so the model can learn negative diagonal entries for the
state transition matrix, which keeps the recurrence stable across long sequences. And the hard
truncation of the convolution output, the [:-d_conv + 1] slice, discards the padding artifacts so
that every position sees exactly the right local window. None of these are visible from the
architecture diagram, but they are precisely the details that make Mamba trainable and stable at
scale.
S-Mamba: Scaling Innovation
Mamba solved the efficiency problem but left a scalability question open: how do you increase model capacity without paying for a full dense expansion? S-Mamba’s answer is to borrow the mixture-of-experts idea and apply it inside the state space block. Instead of one large Mamba block, S-Mamba maintains several expert blocks in parallel, each a full Mamba-style SSM, and learns to route tokens to the experts that are most useful for them. This gives the model more capacity than any single block while activating only a subset of parameters per token, which keeps inference cost close to a much smaller model.
The gating network and the two supporting modules in the code below are the heart of the design. The gate maps the token to a softmax distribution over experts, weighting how much each expert contributes to the final output. The state adapter is a novel addition: a small recurrent cell that produces the state-space state in a modular, learnable way rather than inheriting it implicitly from the block. The fusion module then combines the per-expert outputs into a single representation, using both the gating weights and a learned fusion projection.
class S_MambaBlock:
"""
S-Mamba: Scalable Selective State Space Model
"""
def __init__(self, d_model, d_state=128, d_conv=4, expand=2, num_experts=4):
super().__init__()
self.d_model = d_model
self.d_state = d_state
self.num_experts = num_experts
# Multi-expert selective mechanism
self.experts = nn.ModuleList([
MambaBlock(d_model, d_state, d_conv, expand)
for _ in range(num_experts)
])
# Gating network for expert selection
self.gate = nn.Linear(d_model, num_experts)
# Modular state update
self.state_adapter = StateAdapter(d_model, d_state)
# Mixture fusion
self.fusion = MixtureFusion(d_model, num_experts)
def forward(self, x):
"""
S-Mamba forward with scalable expert selection
"""
# Gate: determine expert weights
gate_weights = F.softmax(self.gate(x), dim=-1)
# Process through experts
expert_outputs = []
for expert in self.experts:
out = expert(x)
expert_outputs.append(out)
# Stack and fuse outputs
expert_tensor = torch.stack(expert_outputs, dim=0) # [num_experts, batch, seq, dim]
# Weighted fusion
fused = self.fusion(expert_tensor, gate_weights)
# Modular state updates
state = self.state_adapter(fused)
return fused, state
class StateAdapter(nn.Module):
"""
Learnable state adapter for modular updates
"""
def __init__(self, d_model, d_state):
super().__init__()
self.state_projection = nn.Linear(d_model, d_state)
self.state_update = nn.GRUCell(d_state, d_state)
def forward(self, x):
"""
Adapt state based on input
"""
state = self.state_projection(x)
# Update state with recurrence
return self.state_update(state)
class MixtureFusion(nn.Module):
"""
Fusion mechanism for combining expert outputs
"""
def __init__(self, d_model, num_experts):
super().__init__()
self.fusion_weights = nn.Linear(d_model * num_experts, num_experts)
self.norm = nn.LayerNorm(d_model)
def forward(self, expert_tensor, gate_weights):
"""
Fuse expert outputs with learnable weights
"""
# expert_tensor: [num_experts, batch, seq, dim]
batch, seq, dim = expert_tensor.shape[1:]
# Flatten experts
flat_experts = expert_tensor.permute(1, 2, 0, 3).reshape(batch, seq, -1)
# Learn fusion weights
fusion_weights = F.softmax(self.fusion_weights(flat_experts), dim=-1)
# Weighted combination
fused = (expert_tensor * fusion_weights.permute(1, 2, 0, 1).unsqueeze(-1)).sum(dim=0)
return self.norm(fused)
The forward pass shows how these pieces fit together at runtime. Gate weights are computed once per
token, then every expert processes the token, and the outputs are stacked and fused by a weighted
combination. From the outside this looks like a dense block, since all experts run, but in a real
deployment with expert pruning the gate can route to only the top-k experts, converting that loop
into a sparse activation. The StateAdapter deserves particular attention because it returns an
explicit state tensor from every block, which the model can then cache across time steps, and the
design uses a GRU cell to update that state recurrently, giving the block a built-in notion of
long-horizon memory that a plain feed-forward expert would lack.
Key Innovations in S-Mamba
1. Input-Conditioned Gating
The first innovation formalizes the gating logic into its own module. Rather than computing gate values from a single projection, input-conditioned gating aggregates the entire input sequence into a fixed-size summary and routes it through a small network that outputs a distribution over gates. Aggregating across the sequence, here with a mean, is what lets the gate make a decision informed by the whole token span rather than a single position, which matters for tasks where the relevant context is spread out.
This module also illustrates a broader design philosophy in S-Mamba: separate concerns. The gating network, the state adapter, and the fusion mechanism are all small, composable modules with clearly defined responsibilities, which makes the architecture easy to extend and easy to reason about. The two-layer structure with a GELU activation in the middle is a standard pattern for a learned softmax gate, and it is flexible enough that the same module can be reused across modalities with no changes.
class InputConditionedGating(nn.Module):
"""
Dynamic gating based on input characteristics
"""
def __init__(self, d_model, num_gates):
super().__init__()
self.gate_network = nn.Sequential(
nn.Linear(d_model, d_model // 2),
nn.GELU(),
nn.Linear(d_model // 2, num_gates),
nn.Softmax(dim=-1)
)
def forward(self, x):
"""
Compute input-dependent gate values
"""
# Aggregate across sequence
x_agg = x.mean(dim=1) # [batch, dim]
# Compute gates
gates = self.gate_network(x_agg)
return gates
2. Structured Parameterization
The second innovation addresses how the state transition matrix, the A matrix that governs how the hidden state evolves, is parameterized. In the original Mamba formulation, A is a full matrix learned as free parameters, which is expressive but expensive and prone to overfitting as the model scales. S-Mamba instead structures A explicitly, allowing a choice between a diagonal matrix, which stores one scalar per state dimension and is far cheaper, and a full matrix for applications that need maximum expressiveness. This switch is controlled by a single boolean flag, which makes it trivial to compare both configurations on the same task.
The trade-off between the two options is the classic expressiveness-versus-efficiency curve. A diagonal A can be seen as modeling each state dimension independently, which is surprisingly powerful because many learned state transitions in practice are nearly diagonal anyway; a full A can capture interactions between state dimensions but scales quadratically in the state size. The learnable projections for B and C are shared between the two modes, so the only structural difference is in A itself. Choosing the diagonal form also pays off at inference time, since the state update becomes a simple element-wise operation that maps directly onto GPU kernels.
class StructuredParameterization(nn.Module):
"""
Structured SSM parameters for better scaling
"""
def __init__(self, d_model, d_state, diagonal=True):
super().__init__()
self.diagonal = diagonal
if diagonal:
# Diagonal A matrix (more efficient)
self.A = nn.Parameter(torch.randn(d_model, d_state))
else:
# Full A matrix (more expressive)
self.A = nn.Parameter(torch.randn(d_model, d_state, d_state))
# Learnable B and C projections
self.B_proj = nn.Linear(d_model, d_state)
self.C_proj = nn.Linear(d_model, d_state)
def forward(self, x):
"""
Compute structured SSM parameters
"""
B = self.B_proj(x)
C = self.C_proj(x)
# Use diagonal A
A = torch.diag(self.A) if self.diagonal else self.A
return A, B, C
3. Parallel Scan Optimization
The recurrent nature of state space models poses a fundamental obstacle to fast training: the output at time t depends on the entire history before it, so a naive sequential implementation would process the sequence one token at a time. Parallel scan algorithms break this dependency by exploiting the associativity of the state space recurrence. Because the operation of applying one state transition followed by another is itself a transition, the recurrence can be re-associated into a tree-like structure, and the whole sequence can be processed in parallel where a sequential loop would be forced to wait on each step.
The implementation below demonstrates the two-phase structure of a practical parallel scan. The sequence is first split into fixed-size chunks, each of which can be scanned independently and concurrently; within each chunk the scan still runs sequentially, but the chunking limits the sequential depth and lets the work be distributed across processing elements. This is a deliberately simplified version of the cooperative scans used in production kernels, which recursively combine partial results, but it captures the essential idea: parallelism comes from re-associating the recurrence, not from removing it.
class ParallelScanSSM:
"""
Efficient parallel scan for SSM computation
"""
@staticmethod
def scan(A, B, C, x):
"""
Parallel scan algorithm for SSM
Computes: y_t = C_t * sum(A_{t-1}...A_0 * B_0 * x_0)
"""
# Chunk for parallel processing
chunk_size = 64
# Compute A powers in chunks
A_chunks = A.chunk(x.size(1) // chunk_size, dim=1)
# Parallel scan within chunks
y_chunks = []
for A_chunk in A_chunks:
y_chunk = S_MambaBlock._parallel_scan(A_chunk, B, x)
y_chunks.append(y_chunk)
# Combine chunks
y = torch.cat(y_chunks, dim=1)
return y
@staticmethod
def _parallel_scan(A, B, x):
"""
Inner parallel scan implementation
"""
# Cooperative scan (simplified)
T = x.size(1)
# Vectorized scan
for i in range(1, T):
x[:, i] = torch.matmul(A[:, i], x[:, i-1]) + B[:, i] * x[:, i]
return x
S-Mamba for Different Modalities
A single architecture that only works for language would be interesting but narrow. The design goal of S-Mamba is modality-agnostic sequence modeling, and the modules built so far are reusable because they all operate on sequences of hidden vectors. Whether those vectors encode tokens, sensor readings, or image patches is irrelevant to the S-Mamba block itself, so the same core can be wrapped in small modality-specific heads. The next three sections show exactly that pattern: the differences between language, time series, and vision models are confined to the input and output projections, while the sequence model in the middle stays unchanged.
Language Modeling
Language modeling is the most demanding application because it combines two hard requirements: linear-time decoding for long prompts and efficient generation of long outputs. The language model below follows the standard transformer-style wrapper around a stack of blocks: an embedding layer maps token ids to vectors, the S-Mamba blocks process the sequence, a final normalization and linear head project to vocabulary logits. The key addition is the explicit return of per-layer states, which the generation loop can cache and feed forward, eliminating the need to reprocess the entire prefix for every new token.
That state caching is what makes autoregressive generation with S-Mamba cheap compared to attention-based models. A transformer must retain a key-value cache whose size grows with context length, while the S-Mamba state is fixed-size regardless of how much history has been seen. When the model returns both logits and states from every layer, the inference loop can pass the states along as it steps through tokens, turning what would otherwise be a quadratic process into a linear one.
class S_MambaLM:
"""
S-Mamba for language modeling
"""
def __init__(self, vocab_size, d_model, num_layers, num_experts=4):
self.embedding = nn.Embedding(vocab_size, d_model)
self.layers = nn.ModuleList([
S_MambaBlock(d_model, num_experts=num_experts)
for _ in range(num_layers)
])
self.norm = nn.LayerNorm(d_model)
self.lm_head = nn.Linear(d_model, vocab_size)
def forward(self, input_ids):
"""
Language modeling forward
"""
x = self.embedding(input_ids)
# Cache states for generation
states = []
for layer in self.layers:
x, state = layer(x)
states.append(state)
x = self.norm(x)
logits = self.lm_head(x)
return logits, states
Time Series Forecasting
Time series forecasting is a natural fit for state space models, and it is the domain where the selective mechanism pays off most visibly. Financial or sensor data is full of regime changes, where the predictive relevance of the distant past varies abruptly, and a model with a fixed-state model either ignores useful history or is swamped by it. Because S-Mamba’s parameters depend on the input, the model can learn to hold onto a trend during a stable period and discard it when the series shifts, which is exactly the behavior a good forecaster needs.
The implementation below uses an encoder-decoder layout: six S-Mamba blocks read the historical window, then two more blocks generate future steps one at a time. The multi-step decoding loop is where the recurrent nature shines. Instead of predicting all future values at once, the model feeds each prediction back into the decoder to produce the next step, compounding its uncertainty exactly as a real forecast should. The output projection collapses the hidden state back to the original feature dimension at each step, and the predictions are concatenated into a single tensor matching the requested horizon.
class S_MambaTimeSeries:
"""
S-Mamba for time series forecasting
"""
def __init__(self, input_dim, d_model, num_experts=4):
self.input_proj = nn.Linear(input_dim, d_model)
self.encoder = nn.ModuleList([
S_MambaBlock(d_model, num_experts=num_experts)
for _ in range(6)
])
self.decoder = nn.ModuleList([
S_MambaBlock(d_model, num_experts=num_experts)
for _ in range(2)
])
self.output_proj = nn.Linear(d_model, input_dim)
def forecast(self, x, horizon):
"""
Multi-step forecasting
"""
# Encode historical data
x = self.input_proj(x)
for layer in self.encoder:
x, _ = layer(x)
# Decode future steps
predictions = []
current = x
for _ in range(horizon):
for layer in self.decoder:
current, _ = layer(current)
pred = self.output_proj(current[:, -1:])
predictions.append(pred)
return torch.cat(predictions, dim=1)
Vision Tasks
Applying a sequence model to images requires first deciding what the sequence is. The vision implementation answers this with a patch embedding, the same trick used by vision transformers: the image is divided into a grid of patches, each patch is flattened and projected into a hidden vector, and the resulting list of vectors becomes a sequence that S-Mamba can process. The grid scanning order matters here, since a linear sequence imposes an ordering on two-dimensional data, and design choices such as scanning rows or diagonals change what spatial structure the model can learn.
The segmentation model below stacks twelve encoder blocks followed by four decoder blocks, a depth
that hints at the scale S-Mamba is designed to support, and ends with a convolutional head that
projects hidden vectors back to per-class logits. The reshape back to spatial dimensions, the H = W = int(N ** 0.5) step, reverses the patch embedding so the output matches the image geometry. What
is notable is that the S-Mamba blocks themselves are identical to the language version, confirming
the claim that one sequence model can serve very different modalities with only the heads changing.
class S_MambaVision:
"""
S-Mamba for vision tasks (image segmentation)
"""
def __init__(self, in_channels, num_classes, d_model=256):
self.patch_embed = PatchEmbed(in_channels, d_model)
self.encoder = nn.ModuleList([
S_MambaBlock(d_model, num_experts=4)
for _ in range(12)
])
self.decoder = nn.ModuleList([
S_MambaBlock(d_model, num_experts=2)
for _ in range(4)
])
self.segmentation_head = nn.Conv2d(d_model, num_classes, 1)
def forward(self, x):
"""
Image segmentation forward
"""
# Convert to patches
x = self.patch_embed(x) # [B, N, D]
# Encode with S-Mamba
for layer in self.encoder:
x, _ = layer(x)
# Decode
for layer in self.decoder:
x, _ = layer(x)
# Reshape to spatial and predict
B, N, D = x.shape
H = W = int(N ** 0.5)
x = x.transpose(1, 2).reshape(B, D, H, W)
return self.segmentation_head(x)
Performance Comparison
Benchmark Results
The benchmark table below is best read as a relative comparison rather than a claim of absolute numbers, since exact results depend heavily on model size, data, and training budget. The consistent story across all three tasks is that S-Mamba improves on both the transformer baseline and the original Mamba. In language modeling, perplexity drops from 15.2 for the transformer to 13.9 for S-Mamba, while inference speed improves to 2.8x relative to the transformer baseline, showing that the accuracy gains come without sacrificing the speed advantage that motivated the state space approach in the first place.
The time series and vision rows tell the same story in different metrics. The mean absolute error on forecasting falls to 0.098, a meaningful step down from Mamba’s 0.128, and mean intersection-over-union on segmentation rises to 82.1. What these numbers suggest is that the mixture-of-experts capacity increase and the explicit state adapter pay off across modalities, not just in language. It is worth keeping in mind that the speed advantage grows with sequence length: the linear-time property is the entire reason S-Mamba can process contexts that would be impractical with attention.
benchmarks = {
'language_modeling': {
'perplexity': {
'Transformer': 15.2,
'Mamba': 14.8,
'S-Mamba': 13.9
},
'inference_speed': {
'Transformer': '1.0x',
'Mamba': '2.1x',
'S-Mamba': '2.8x'
}
},
'time_series': {
'mae': {
'Transformer': 0.142,
'Mamba': 0.128,
'S-Mamba': 0.098
}
},
'vision_segmentation': {
'mIoU': {
'Transformer': 78.5,
'Mamba': 79.2,
'S-Mamba': 82.1
}
}
}
Memory Efficiency
Performance per se is only half the argument for S-Mamba; memory efficiency is the other half, and for many deployment scenarios it is the deciding factor. The comparison below isolates the memory story, showing that all three models are trained at the same 7B parameter count, so the differences come entirely from how each architecture manages its runtime state rather than from model size. The theoretical key-value cache complexity tells the story: a transformer’s cache grows quadratically with sequence length N, Mamba’s grows with N times the state dimension, and S-Mamba’s is independent of sequence length thanks to the fixed-size state adapter.
memory_comparison = {
'parameters': {
'Transformer_7B': '7B',
'Mamba_7B': '7B',
'S-Mamba_7B': '7B'
},
'kv_cache': {
'Transformer': 'O(N²)',
'Mamba': 'O(N × d_state)',
'S-Mamba': 'O(d_state)'
},
'inference_memory_8k': {
'Transformer': '48GB',
'Mamba': '24GB',
'S-Mamba': '18GB'
}
}
The concrete numbers for an 8,000-token context illustrate the practical impact of those complexity classes. Serving the transformer requires 48GB of memory for its key-value cache, Mamba roughly halves that at 24GB, and S-Mamba comes in at 18GB. This is the difference between fitting a model on a single GPU and spilling across multiple devices, and it translates directly into cost per inference. The trade-off to weigh is that this efficiency is purchased with some of the architectural complexity introduced earlier: expert routing, the state adapter, and structured parameterization all add implementation overhead that only pays off at longer sequences or higher throughput.
Implementation Best Practices
When to Use S-Mamba
Knowing what an architecture is good at is only useful if you also know when to walk away from it. The checklist below encodes the conditions under which S-Mamba’s design choices win. The strongest signals are long sequences, where linear complexity becomes a decisive advantage, and memory-constrained serving, where the smaller state footprint matters. Multi-modal workloads benefit because one unified architecture can share a single code path across domains, and real-time applications gain from the fast autoregressive decoding.
use_s_mamba_when = {
'long_sequences': True, # Linear complexity is key
'limited_memory': True, # Smaller KV cache
'multi_modal': True, # Unified architecture
'real_time': True, # Fast inference needed
'not_ideal_for': [
'short_sequences', # Overhead not worth it
'simple_tasks', # Simpler models suffice
]
}
The explicit counter-indications are just as important as the positive signals. For short sequences, the expert routing and state machinery add overhead with little benefit, since a transformer at that scale is perfectly affordable and simpler to reason about. For simple tasks, a smaller feed-forward or attention model will train faster and be easier to debug. The discipline of defining these boundaries up front is what separates a thoughtful adoption of S-Mamba from a fashionable but inappropriate one, and the same reasoning applies to any efficient-sequence-model decision.
Conclusion
S-Mamba represents a significant advancement in state space models:
- Scalability: Modular expert selection enables scaling
- Efficiency: Linear complexity with smaller memory footprint
- Versatility: Works across language, vision, and time series
- Performance: Outperforms both Transformer and Mamba in benchmarks
As research continues, S-Mamba and similar architectures may become the foundation for next-generation efficient AI systems.
Comments