Mamba-3: Improved Sequence Modeling using State Space Principles

A compact memory needs a better update rule

A recurrent language model carries a bounded state from one token to the next. That makes decoding attractive: the model can update its memory without revisiting an ever-growing cache of past keys and values. But bounded memory introduces two separate questions. Can the update represent the computation a task requires? And does the update use the GPU well enough for the theoretical savings to matter?

Mamba-3 addresses both questions by returning to the continuous-time state space model behind the recurrence. It changes how inputs enter the state, restores rotational dynamics for state tracking, and performs more useful computation per state load through multi-input, multi-output updates. The result is an evolution of Mamba-2’s structured state space duality, rather than a replacement of recurrence with attention. The paper is by Lahoti and colleagues, listed as ICLR 2026 in its versioned arXiv record.

The distinction matters when interpreting the headline. A stronger fixed-size memory is still fixed-size memory. Mamba-3 improves the reported language-modeling tradeoff and synthetic state tracking, while retaining weaknesses on some retrieval tasks. We will derive the changes first, then separate what the experiments demonstrate from what the architecture merely makes possible.

Start with Mamba-2, then change the input integral

Write the continuous system as a state derivative and a readout:

Mamba-2 restricts the transition generator to a scalar times identity within a head. The discrete transition is therefore a positive scalar decay, while input-dependent projections determine what to write and read. In the notation used by this paper, its implemented update is:

This is exponential-Euler, not the exact zero-order-hold input coefficient. The exponential handles the state transition; the input integral uses a right-endpoint approximation. For fixed parameters, exact ZOH would instead use A^{-1}(e^{\Delta A}-I)B x. The new paper formalizes the earlier implementation through its exponential-adjusted discretization framework. §2.2 and §3.1.1

After factoring out the exponential dynamics, Mamba-3 approximates the remaining input integral using both endpoints. Let v_t=B_tx_t. Then:

The previous input is not simply concatenated to the current token. It is a previous state-input, already containing the interaction between B and x. The learned scalar \lambda_t\in[0,1] controls the endpoint mixture. Setting it to one recovers exponential-Euler; setting it to one half gives the symmetric exponential-trapezoidal approximation. Proposition 1 and Remark 2

There is a numerical-analysis qualification here. The second-order input-integral error statement requires \lambda_t=\tfrac12+O(\Delta_t) under the stated assumptions. The authors do not enforce this condition in their best empirical model. A free learned endpoint weight should therefore be understood as an expressivity choice; “trapezoidal” does not guarantee second-order accuracy for arbitrary learned parameters. Remark 3 and Appendix A

The extra endpoint is an implicit local filter

The recurrence first combines two adjacent state-inputs, then sends the result through the decaying memory. In matrix form, this becomes a decay operator multiplied by a lower two-band operator. Mamba-2 has only the diagonal input scaling in that second factor. Mamba-3 gains a subdiagonal, which sends an input into the next step as well as its own step. §3.1.2–3.1.3, Eq. 7

This is related to convolution, but its location matters. An external short convolution acts on x, and often B and C, before the recurrence. Here the local mixture acts on Bx inside the recurrence. Together with B/C biases, this lets the authors remove the external short causal convolution and its activation in their tested architecture. It does not establish that every recurrent network can discard its local preprocessing. Remark 4 and §3.4

Decay remembers magnitude; rotation can remember state

A scalar decay can attenuate an old state, but cannot rotate a two-dimensional state around a circle. That is a meaningful limitation for a task such as parity. One exact construction starts at (1,0), rotates by π when a one arrives, and does nothing for zero. The sign of the horizontal coordinate then records whether the prefix has an even or odd number of ones.

Mamba-3 introduces a complex transition generator. In real coordinates, each complex entry becomes a two-dimensional pair with a decayed rotation:

A complex state of size N/2 corresponds to N real coordinates. Counting a complex value as if it were one real scalar would create a misleading memory comparison. The readout equivalence also includes a sign convention: the real representation of the complex C projection negates its imaginary component. Proposition 2

Keep the scalar recurrence by changing coordinates

Explicitly rotating every state pair is not the only implementation. Define the accumulated phase \Phi_t=\sum_{i=0}^t\phi_i and express the state in a rotating coordinate frame. The transition becomes scalar again, while B and C receive rotations by the negative cumulative phase. This is the paper’s data-dependent RoPE equivalence. Unlike ordinary positional RoPE, these increments depend on the input. Proposition 3

The delayed endpoint needs special care. Its B projection belongs to the previous accumulated phase, while the current write belongs to the current phase. Rotating both endpoints with the current phase would silently change the recurrence. Proposition 4 combines these frame changes with exponential-trapezoidal discretization. That combined identity is one of the checks in the reference implementation below. Eq. 11

These constructions establish representational capability, not automatic learnability. The paper supplies separate training evidence: its formal-language table reports scaled scores of 100.00 for parity, 98.51 for arithmetic without brackets and 87.75 with brackets. Those are scaled accuracies, not ordinary raw classification accuracy. Standard RoPE and no-RoPE ablations perform poorly. GDN with eigenvalues allowed in [-1,1] also solves parity and scores higher on the two arithmetic tasks, so state tracking is not unique to complex SSMs. Table 5(b)

MIMO spends arithmetic on the memory already being loaded

A single head has a matrix state H\in\mathbb{R}^{N\times P}. In the SISO formulation, B is a length-N vector and x is a length-P vector; their outer product writes into that matrix. The state is large enough that moving it can dominate the small amount of arithmetic used to update it.

MIMO gives B and x a rank dimension R:

The state remains N×P; the write is now a matrix product, and C produces R readout channels. This simplified two-term equation isolates the rank change, just as §3.3 does. The deployed Mamba-3 recurrence also contains the delayed input and rotations. §3.3, Eqs. 12–14

The inference argument is about arithmetic intensity. For small R relative to N and P, the dominant state traffic changes little while the write/read computation grows. A memory-bound kernel may use previously idle compute capacity. That creates room for better quality at similar decode latency; it is not a universal promise that increasing R costs nothing.

Training has a different tradeoff. A naive black-box decomposition has R² SISO interactions. The chunked algorithm reduces the chunk length by R, making the analyzed total overhead linear in R when the dimensions are balanced. Actual runtime still depends on kernels and dimensions. The architecture also avoids multiplying all dense projections by R: it uses learned elementwise rank expansion for x, y and the gate, and reduces the MLP width to match parameter counts. At 1.5B, the reported MLP inner width changes from 4096 to 3824. §3.3 and §4.1.1

The block contains more than the three headline ideas

B and C receive RMS normalization followed by learned head-specific channel biases. The paper associates these biases and the two-endpoint update with useful convolution-like behavior. Its 440M component ablation reports test perplexity of 16.68 without bias and trapezoidal updates, 16.49 without bias, 15.72 for the full model and 15.85 after adding a short convolution. This is evidence for the combined design in that experiment, not a universal causal ranking of all local filters. §3.4 and Table 5(a)

Normalization placement is also consequential. Pure Mamba-3 removes Mamba-2’s post-gate RMSNorm after adding BCNorm. In hybrids, reintroducing pre-gate grouped RMSNorm improves length generalization in the reported retrieval experiment, with other tradeoffs. The paper explicitly leaves the ideal norm type and placement unresolved. Treating the pure-model diagram as a universal hybrid recipe would miss that result. §4.1.2 and Appendix E

Read the quality and latency evidence separately

The main language-model comparison trains all models on 100B FineWeb-Edu tokens, using the Llama-3.1 tokenizer and 2K context. At 1.5B, SISO Mamba-3 averages 56.4% on the downstream suite versus 55.8% for Gated DeltaNet. Rank-four MIMO reaches 57.6%, another 1.2 percentage points. FineWeb-Edu perplexity improves from 10.47 for Mamba-2 to 10.35 for SISO and 10.24 for MIMO. These are matched-protocol results, not comparisons against arbitrary released checkpoints. Table 3 and Appendix D

The retrieval results qualify the language-model average. On SWDE, pure Mamba-3 SISO scores 28.5 and MIMO 36.3, while the Transformer scores 48.9. On FDA, MIMO scores 29.3 versus the Transformer’s 58.4. An improved compressed state does not provide the same ability to revisit individual tokens as an attention cache. The hybrid experiments use five linear layers per attention layer and NoPE attention, so they should not be conflated with the pure recurrent results. Table 4

“Similar latency” has a workload and a boundary

Table 6 measures decode-kernel latency for one step at batch 128 on one H100, using dimensions from 1.5B models: model width 2048, state size 64 or 128, and FP32 or BF16. These kernel measurements are not complete language-model token latency. At BF16 and state size 128, the reported times are 0.203 ms for Mamba-2, 0.156 ms for Mamba-3 SISO and 0.179 ms for MIMO. At state size 64, MIMO is 0.137 ms versus Mamba-2’s 0.127 ms. MIMO is therefore not uniformly faster at every setting. These are the authors’ measurements; no timing was collected for this review. Table 6 and §4.4

The paper compares its CuTe DSL decode implementation with released Triton baselines, so the observed speed combines architecture and kernel engineering. Table 7 separately reports prefill and prefill-plus-decode wall-clock times in seconds, as specified in Appendix G.2. For a 2048-token prefill, MIMO takes 2.42 s versus SISO’s 2.02 s. This benchmark uses batch 128 on an H100-SXM 80GB and three repetitions; the combined workload additionally decodes 2048 tokens per sequence. These whole-workload times must not be read in the millisecond units of Table 6. They expose a prefill cost hidden by a decode-only headline. Figure 3’s state-size sweep uses 440M models and a different token-budget protocol; its smaller-state quality result should not be substituted directly for the main 1.5B experiment. §4.3–4.4, Appendix G.2: timing units and workload

A useful interpretation is that Mamba-3 offers a better set of options for designing a recurrent mixer. It expands the update’s expressivity, restores an important kind of state transition, and provides a hardware-motivated way to trade more arithmetic for better quality. The unresolved questions concern retrieval under bounded memory, hybrid normalization, and how much of the latency advantage survives a different serving workload.

Five executable views of the equations

The panels below are extracted from examples/mamba3.py. They execute the arithmetic discussed here using the Python standard library. They do not train a language model or load weights. Official source links are pinned to e9594ce.

1. Keep the delayed endpoint inside the recurrence

Original paper

Enlarge figure

Math and tensor shapes

\begin{aligned}h_t&=\alpha_t h_{t-1}+\beta_t v_{t-1}\\&\quad+\gamma_t v_t,\\v_t&=B_tx_t\end{aligned}

The scalar helper takes fixed Δ, A and λ, with h[-1]=v[-1]=0. β=(1−λ)Δ exp(ΔA), γ=λΔ. λ=1 removes the delayed write and recovers the Mamba-2 exponential-Euler update. It is not the exact ZOH input coefficient. Learned parameters would vary by token and head.

inputs [T] → states [T]

Inspect official block and kernel calls

Python implementation

python

def coefficients(delta, a, weight):
    """Returns alpha, beta, gamma for the exponential-trapezoidal rule.

    Args:
        delta: Nonnegative time interval.
        a: Nonpositive real decay generator.
        weight: Current-endpoint weight in [0, 1].

    Returns:
        Three scalar coefficients from paper Eq. 6.
    """
    if not all(map(math.isfinite, (delta, a, weight))):
        raise ValueError("Finite parameters required")
    if delta < 0 or a > 0 or not 0 <= weight <= 1:
        raise ValueError("Require delta >= 0, a <= 0, weight in [0, 1]")
    alpha = math.exp(delta * a)
    return alpha, (1 - weight) * delta * alpha, weight * delta


def scan(inputs, delta, a, weight):
    """Scans scalar state-inputs v=B*x, with h[-1]=v[-1]=0."""
    alpha, beta, gamma = coefficients(delta, a, weight)
    state, previous = 0.0, 0.0
    outputs = []
    for value in inputs:
        state = alpha * state + beta * previous + gamma * value
        outputs.append(state)
        previous = value
    return outputs

2. Expand the causal operator independently

Original paper

Enlarge figure

Math and tensor shapes

\begin{aligned}L_{ts}&=\gamma\alpha^{t-s}\\&\quad+\mathbf{1}_{s<t}\beta\alpha^{t-s-1},\\&\qquad s\le t\end{aligned}

For constant coefficients and B=C=1, a token writes immediately with γ and once more at the next step with β. Both contributions then decay. This helper expands the operator without calling the sequential scan. Multiplying it by the input is an independent equivalence check, not an implementation of a fast chunked kernel.

length T → causal matrix [T,T]

Inspect official block and kernel calls

Python implementation

python

def causal_mask(length, delta, a, weight):
    """Expands the recurrence into a lower triangular scalar input map."""
    if not isinstance(length, int) or length < 1:
        raise ValueError("Positive length required")
    alpha, beta, gamma = coefficients(delta, a, weight)
    return [
        [0.0 if s > t else gamma * alpha ** (t - s)
         + (beta * alpha ** (t - s - 1) if s < t else 0.0)
         for s in range(length)]
        for t in range(length)
    ]

3. Move rotations out of the state update

Original paper

Enlarge figure

Math and tensor shapes

\begin{aligned}\Phi_t&=\sum_{i=0}^{t}\phi_i,\\\widetilde B_t&=R(-\Phi_t)B_t,\\\widetilde C_t&=R(-\Phi_t)C_t\end{aligned}

The physical frame rotates both the retained state and the previous input write before adding the new write. The transformed frame rotates the current write by −Φt, retains the already transformed previous write, and unrotates the state for readout. Confusing the previous and current phases breaks the trapezoidal equivalence. One real pair stands in for the block diagonal state; B=C=(1,0).

inputs [T], angles [T] → real outputs [T]

Inspect official block and kernel calls

Python implementation

python

def rotate(vector, angle):
    """Applies one real 2-by-2 rotation, equivalent to complex phase."""
    x, y = vector
    return (math.cos(angle) * x - math.sin(angle) * y,
            math.sin(angle) * x + math.cos(angle) * y)


def rotating_scan(inputs, angles, delta, a, weight, transformed=False):
    """Computes physical or Eq. 11 transformed-frame outputs for B=C=(1,0)."""
    if len(inputs) != len(angles):
        raise ValueError("One angle per input required")
    alpha, beta, gamma = coefficients(delta, a, weight)
    state, previous, phase = (0.0, 0.0), (0.0, 0.0), 0.0
    outputs = []
    for value, angle in zip(inputs, angles):
        phase += angle
        write = rotate((value, 0), -phase) if transformed else (value, 0)
        retained = state if transformed else rotate(state, angle)
        prior = previous if transformed else rotate(previous, angle)
        state = tuple(alpha * h + beta * old + gamma * new
                      for h, old, new in zip(retained, prior, write))
        previous = write
        physical = rotate(state, phase) if transformed else state
        outputs.append(physical[0])
    return outputs

4. Sum rank-one writes without expanding the recurrent state

Original paper

Enlarge figure

Math and tensor shapes

\begin{aligned}W_t&=B_tX_t^\top,\\H_t&=\alpha_tH_{t-1}+\Delta_t W_t,\\Y_t&=H_t^\top C_t\end{aligned}

B has shape N×R, X has shape P×R, H remains N×P and the readout has shape P×R. This row isolates the write. The full block additionally uses the two-endpoint rule, rotations, learned rank expansion, gating and output combination. Rank one recovers an outer product exactly.

B [N,R], X [P,R] → write [N,P]

Inspect official block and kernel calls

Python implementation

python

def mimo_write(b, x):
    """Computes B @ X.T with shapes [N,R], [P,R] -> [N,P]."""
    if not b or not x or not b[0]:
        raise ValueError("Nonempty matrices required")
    rank = len(b[0])
    if any(len(row) != rank for row in b + x):
        raise ValueError("Matching rectangular rank required")
    return [[sum(br * xr for br, xr in zip(row, column))
             for column in x] for row in b]

5. Count arithmetic before making a latency claim

Original paper

Enlarge figure

Math and tensor shapes

I_R=\frac{4NPR+NP-PR}{2(1+2NR+PR+NP)}

This is the paper’s simplified 2-byte accounting convention. At R=1 the numerator becomes 5NP−P, the SISO count. State elements remain NP while projected inputs grow with R. Actual traffic depends on state read/write, fusion, cache behavior and dtype. This value neither estimates end-to-end throughput nor proves MIMO is free.

positive integer N,P,R → modeled FLOPs/byte

Inspect official block and kernel calls

Python implementation

python

def arithmetic_intensity(n, p, rank):
    """Returns Table 2's modeled FLOPs/byte for 2-byte values, not latency."""
    if any(not isinstance(v, int) or v < 1 for v in (n, p, rank)):
        raise ValueError("Positive integer dimensions required")
    flops = 4 * n * p * rank + n * p - p * rank
    byte_count = 2 * (1 + 2 * n * rank + p * rank + n * p)
    return flops / byte_count

Original paper excerpts: Aakash Lahoti, Kevin Y. Li, Berlin Chen, Caitlin Wang, Aviv Bick, J. Zico Kolter, Tri Dao, Albert Gu. Paper v1. Cropped/rasterized from the original PDF; no scientific content redrawn.

Download Python referenceStandard-library requirements

Run locally: python3 examples/mamba3.py --smoke. The JSON contains calculated values, not checkpoint predictions. See examples/README-mamba3.md for assumptions and omitted components.

Related reading

• Mamba introduces selection: token-dependent parameters determine which information enters and survives in state.
• Mamba-2 develops state space duality and the scalar-transition machinery that Mamba-3 extends.
• VMamba shows how scan order and spatial structure become separate design decisions when recurrence moves to images.
• Linear attention approximations connects compressed recurrent state with the attention viewpoint.
• The original Mamba-3 paper, v1 provides derivations, formal-language protocols and kernel benchmark details.
• Official Mamba-3 block at revision e9594ce connects the mathematical parameters to projection layout, BC normalization, biases and kernel dispatch. This code snapshot is later than the paper; its defaults are not evidence of the exact training configuration used for every table.