State-Space Models Long-Context Sequence Architectures Mamba

# State-Space Models & Long-Context Sequence Architectures (Mamba)

## Introduction & Motivation

Transformers have dominated sequence modeling since 2017, but their self-attention mechanism carries a fundamental scaling cost: computing attention over a sequence of length L requires comparing every position to every other position, giving quadratic time and memory complexity in sequence length. This quadratic cost becomes prohibitive for very long sequences, whether long documents, high-resolution audio, genomic sequences, or extended multi-turn conversations, motivating a resurgence of interest in alternative sequence architectures that scale sub-quadratically, or even linearly, with sequence length.

Structured state-space models (SSMs) re-adapt a classical tool from control theory and signal processing, the linear state-space representation of a dynamical system, into a learnable neural sequence layer. The S4 model (Structured State Space Sequence model) demonstrated that a carefully parameterized, initialized, and discretized continuous-time state-space model could be trained efficiently on long sequences and achieve strong results on tasks specifically designed to test long-range dependency modeling, dramatically outperforming Transformers on some of these benchmarks at the time of its introduction.

Mamba, introduced in 2023, extended this line of work with "selective" state-space models, in which the model's dynamics (the parameters governing how the hidden state evolves and how it is read out) are allowed to depend on the input at each time step, rather than being fixed (input-independent) as in earlier linear SSMs such as S4. This selectivity closes much of the expressive gap between SSMs and attention-based models (since a fixed, input-independent linear recurrence cannot, for example, decide to selectively "remember" or "forget" specific tokens based on their content) while preserving a recurrent formulation that allows constant per-step inference cost and linear-time sequence processing, in contrast to the quadratic cost of attention.

The practical significance of this line of research lies in the promise of efficient long-context modeling: an SSM-based language model can, in principle, process arbitrarily long sequences with constant memory per generated token (since it maintains a fixed-size hidden state rather than an ever-growing cache of past keys and values as in Transformer attention), and Mamba-based and hybrid Mamba-Transformer architectures have demonstrated competitive language modeling quality with substantially improved inference throughput and memory efficiency at long context lengths, making them an active and closely watched alternative and complement to the standard Transformer architecture.

## Core Concepts & Theory

A state-space model, in its classical continuous-time form, describes how a hidden state vector evolves over time in response to an input signal, and how an output is read out from that hidden state. The state evolves according to a linear ordinary differential equation driven by the input, and the output is a linear readout of the current state (plus, often, a direct pass-through term from the input to the output). This is precisely the same mathematical object used to describe linear dynamical systems in classical control theory, such as the state of a physical system evolving under a set of forces.

To use this continuous-time formulation as a layer within a neural network operating on discrete sequences (such as tokens or discretized audio samples), the continuous-time system must be discretized, converting the continuous differential equation into a discrete-time recurrence that can be applied step by step across the elements of a sequence, using a discretization method (commonly a zero-order hold or a bilinear transform) parameterized by a learnable or input-dependent step size.

The key architectural insight of S4 was that naively parameterizing and training such a discretized linear recurrence performs poorly at capturing long-range dependencies, and that a specific structured initialization of the state matrix, based on the HiPPO (High-Order Polynomial Projection Operators) framework, dramatically improves the model's ability to compress and retain information about long input histories within a fixed-size hidden state. HiPPO initialization sets the state matrix so that the hidden state evolves to maintain an online, compressed approximation of the entire input history projected onto a basis of orthogonal polynomials, giving the model a principled, mathematically motivated starting point for long-range memory rather than relying on gradient descent to discover such a structure from a generic random initialization.

Mamba's central innovation, selectivity, modifies this picture by making the discretization step size, along with certain other recurrence parameters, functions of the current input token rather than fixed, input-independent constants shared across the whole sequence. This allows the model to dynamically control how much the hidden state changes at each step: a large effective step size lets a token's influence on the state decay quickly (effectively "forgetting" less relevant content), while a small effective step size lets a token's information persist and propagate further into the future, a form of content-based, input-dependent gating that is conceptually similar to the gating mechanisms in LSTMs, but derived from and integrated within the state-space formalism.

## Mathematical Formulation

The continuous-time linear state-space model describing the evolution of a hidden state h(t) in response to a scalar input signal x(t), and the readout of a scalar output y(t), is given by a pair of linear equations:

$$ h'(t) = A h(t) + B x(t), \qquad y(t) = C h(t) + D x(t) $$

where A is the state transition matrix governing the unforced dynamics of the hidden state, B maps the input into the state space, C reads the output from the current state, and D is a direct feedthrough term from input to output (commonly omitted or treated as a simple residual connection in practice).

Discretizing this continuous-time system with a step size delta (commonly denoted with the Greek letter capital delta), using a zero-order hold discretization, yields discrete-time recurrence matrices A-bar and B-bar computed from the continuous-time A and B matrices and the step size:

$$ \bar{A} = \exp(\Delta A), \qquad \bar{B} = (\Delta A)^{-1} (\exp(\Delta A) - I) \cdot \Delta B $$

which then define a discrete-time linear recurrence over the sequence of discrete time steps (tokens), applied identically at every position when delta, A, B, and C are shared, fixed parameters, as in the original S4 formulation:

$$ h_t = \bar{A} h_{t-1} + \bar{B} x_t, \qquad y_t = C h_t $$

In Mamba's selective formulation, the parameters B, C, and the step size delta are no longer fixed but are instead computed as functions of the current input token x_t, typically via small learned linear projections applied to x_t:

$$ \Delta_t = au( ext{Linear}_\Delta(x_t)), \qquad B_t = ext{Linear}_B(x_t), \qquad C_t = ext{Linear}_C(x_t) $$

where tau is a nonlinearity (softplus is used in the original Mamba formulation) ensuring the effective step size remains positive, and the discrete-time recurrence is then applied with these input-dependent, time-varying parameters at each step rather than a single shared set of parameters for the whole sequence, which is what gives the model its content-aware, selective behavior.

## Advanced Theory & Extensions

Because the selective recurrence uses input-dependent parameters at every time step, it can no longer be computed via the same global convolution formulation used to accelerate training of the original, non-selective S4 model (where a fixed linear recurrence over a sequence can be equivalently expressed as, and computed via, a single large convolution across the entire sequence, which is highly parallelizable using fast Fourier transform techniques). Mamba instead relies on a custom hardware-aware parallel scan algorithm, implemented with careful attention to GPU memory hierarchy (keeping intermediate states in fast on-chip SRAM rather than repeatedly reading and writing to slower high-bandwidth memory), to compute the input-dependent recurrence efficiently in parallel during training, despite the recurrence no longer having the simple, input-independent convolutional structure that made the original S4 model straightforward to parallelize.

Mamba-2 reformulated the selective SSM recurrence to reveal a mathematical duality with a particular structured form of linear attention, showing that selective state-space models and certain linear-attention variants can be understood as two different computational realizations of a closely related underlying operation (termed "structured state-space duality" by the authors), a connection that helped unify what had previously appeared to be two largely separate architectural research threads (SSMs and linear attention) and enabled further efficiency improvements by borrowing optimized computational techniques developed for each.

Hybrid architectures that interleave SSM layers (or Mamba blocks) with a smaller number of standard attention layers within the same network have become an increasingly common design pattern, motivated by empirical findings that pure SSM architectures, despite their efficiency advantages, sometimes underperform pure Transformers on tasks requiring precise retrieval of specific information from earlier in the context (such as exact copying or in-context "needle in a haystack" retrieval tasks), a shortcoming that a modest number of interspersed attention layers appears to substantially mitigate while retaining most of the efficiency benefits of the SSM backbone for the majority of the network's layers.

## Computational Considerations

The principal computational advantage of SSM-based sequence models over standard Transformer attention is the shift from quadratic to linear time complexity in sequence length during training (since the recurrence, computed via a parallel scan, processes the sequence in time roughly proportional to sequence length rather than sequence length squared), and, even more significantly, a shift from a linearly growing memory footprint per generated token during autoregressive inference (the Transformer's key-value cache, which must store a representation for every previously generated token) to a constant memory footprint per generated token (the SSM's fixed-size hidden state, which does not grow with sequence length).

This constant-memory-per-token inference property is particularly valuable for very long context or long-generation scenarios, where a Transformer's KV cache can come to dominate total memory usage and become the binding constraint on batch size or maximum context length servable on a given accelerator, whereas an SSM's hidden state size is fixed by the model's architecture regardless of how long the generated or processed sequence becomes.

The hardware-aware parallel scan implementation used by Mamba is essential to realizing these theoretical efficiency advantages in practice on real GPU hardware, since a naive, sequential implementation of the recurrence (looping over time steps one at a time) would be far slower in wall-clock training time than a well-optimized attention implementation, despite having a lower asymptotic complexity; the practical speed advantage of SSMs over attention at moderate sequence lengths depends substantially on this careful, hardware-conscious implementation work, not merely on the underlying mathematical formulation.

## Practical Implementation Strategies

Implementing a state-space sequence layer from scratch requires careful attention to numerical stability during discretization, particularly ensuring the discretized state transition matrix A-bar does not produce an unstable (exploding) recurrence, which is typically addressed by initializing and constraining the continuous-time A matrix (for instance, restricting it to have eigenvalues with negative real parts, corresponding to a stable, non-diverging underlying continuous-time dynamical system) and by using the HiPPO-motivated structured initialization discussed above rather than a generic random initialization.

For practitioners building on top of existing implementations rather than implementing the parallel scan from scratch, using established, hardware-optimized SSM libraries (rather than reimplementing the custom CUDA kernels used in the original Mamba release) is strongly preferable, since the wall-clock performance advantage of these architectures depends heavily on low-level implementation details that are easy to get wrong and difficult to match without significant systems engineering effort.

When designing a hybrid Mamba-attention architecture, a common practical starting point is to replace the large majority (roughly 80 to 90 percent) of a Transformer's attention layers with SSM blocks, retaining a small number of full attention layers interspersed at regular intervals, and then empirically evaluating on both efficiency metrics (throughput, memory usage at long context) and quality metrics (particularly targeted long-context retrieval evaluations, which are more sensitive to this architectural trade-off than aggregate language modeling perplexity) to tune the specific ratio for a given application.

## Benchmark Datasets & Evaluation

The Long Range Arena (LRA) benchmark suite, introduced specifically to evaluate long-sequence modeling architectures, includes tasks such as classifying long text sequences, long list operations, image classification on flattened pixel sequences, and pathfinder tasks requiring the model to trace long-range visual dependencies, all constructed with sequence lengths and dependency structures designed to stress-test an architecture's ability to model long-range relationships, and S4's strong performance on this benchmark, relative to Transformers and prior recurrent architectures, was a key early result establishing SSMs as a serious long-sequence modeling alternative.

Standard large language model benchmarks, including perplexity on held-out text corpora and downstream task accuracy on suites such as MMLU, HellaSwag, and ARC, are used to compare Mamba and hybrid Mamba-Transformer language models against comparably-sized Transformer baselines, generally showing that Mamba-class models can match Transformer quality at equivalent parameter and training-compute budgets on many standard tasks.

Targeted long-context evaluations, particularly "needle in a haystack" style tasks (where a specific piece of information is inserted at a controlled position within a long context and the model is tested on its ability to retrieve that information accurately) and synthetic associative recall tasks, have proven especially important for evaluating SSM-based architectures specifically, since these targeted retrieval-style evaluations reveal architectural weaknesses (particularly in pure SSM models without any attention layers) that aggregate perplexity metrics tend to obscure.

## Key Challenges & Limitations

Pure SSM architectures, despite strong performance on many aggregate language modeling and long-range dependency benchmarks, have shown measurable weaknesses on tasks requiring precise, exact retrieval or copying of specific tokens from earlier in a long context, a limitation plausibly connected to the fixed-size hidden state acting as an information bottleneck that must compress the entire input history, in contrast to attention's ability to directly and exactly access any specific prior token's representation without any compression loss.

The training and inference efficiency advantages of SSM architectures depend heavily on specialized, hardware-aware implementations (custom parallel scan kernels tuned to specific GPU memory hierarchies), meaning the practical performance gap between a well-optimized SSM implementation and a well-optimized attention implementation (such as FlashAttention, which similarly uses hardware-aware techniques to reduce attention's memory-access overhead, though it does not change attention's underlying quadratic compute scaling) can be smaller in practice than the difference in asymptotic complexity alone would suggest, particularly at the moderate sequence lengths common in many current applications.

The relative immaturity of the SSM and Mamba software and tooling ecosystem, compared to the many years of accumulated optimization, tooling, and community familiarity built around the Transformer architecture, represents a practical adoption barrier, including less mature support in some standard training and inference serving frameworks, fewer pretrained model checkpoints available across different scales and domains, and less accumulated practitioner experience with SSM-specific training instabilities and hyperparameter sensitivities relative to the extensively documented Transformer training playbook.

## Hyperparameter Tuning

The dimensionality of the hidden state (the size of the state vector h) is a central architectural hyperparameter controlling the model's memory capacity, with larger state dimensions increasing the amount of information the recurrence can retain about the input history at the cost of increased compute and memory per layer, analogous to the role of the key-value cache dimensionality in an attention layer, though the SSM's state size remains fixed regardless of sequence length rather than growing with it.

The initialization scheme for the continuous-time state matrix A, whether using the HiPPO-based structured initialization or a simpler alternative, has a substantial empirical effect on training stability and the model's ability to capture long-range dependencies, and ablation studies in the S4 and Mamba literature consistently show that departing from a well-motivated structured initialization toward a naive random initialization measurably degrades long-range performance, making this a hyperparameter choice with outsized importance relative to its modest apparent complexity.

The ratio and placement of attention layers within a hybrid Mamba-attention architecture, discussed above as a practical implementation consideration, also functions as a tunable hyperparameter trading off efficiency against retrieval-sensitive task quality, and different published hybrid architectures have explored different ratios and placement strategies (uniformly interspersed versus concentrated at specific depths) without full convergence yet on a single clearly optimal configuration across all model scales and tasks.

## Real-World Applications & Case Studies

Long-document and long-context language modeling applications, including processing entire books, lengthy legal or financial documents, or extended multi-turn conversation histories, are a natural fit for SSM and hybrid architectures given their favorable memory scaling at long sequence lengths, and several production and research language models have adopted Mamba or hybrid Mamba-attention backbones specifically to improve long-context serving efficiency relative to pure Transformer alternatives.

Genomic and DNA sequence modeling has emerged as an application area particularly well-suited to SSM architectures, since genomic sequences can span hundreds of thousands to millions of base pairs, far exceeding the practical context lengths of standard attention-based models, and specialized genomic foundation models built on SSM backbones have demonstrated the ability to model such extremely long biological sequences directly, without the aggressive chunking or subsampling strategies that would otherwise be required to fit within a quadratic-attention model's practical context window.

Audio and time-series modeling, including raw waveform audio generation and processing (where sequences can span tens of thousands of samples per second of audio) and long-horizon time-series forecasting, have also been active application areas for SSM architectures, building on the original S4 model's strong results on related long-range sequence benchmarks and the natural fit between the state-space formalism's origins in continuous-time dynamical systems modeling and the continuous, temporally structured nature of audio and sensor time-series data.

## Integration with Other Methods

SSM layers are frequently combined with standard Transformer components beyond just attention, including layer normalization, feedforward (MLP) blocks, and standard token embedding and positional handling schemes (with many SSM architectures, including Mamba, notably not requiring explicit positional embeddings at all, since the recurrent formulation inherently processes tokens in sequential order, unlike the permutation-invariant self-attention operation, which requires explicit positional information to be added).

The structured state-space duality identified in Mamba-2, connecting selective SSMs to a class of linear attention mechanisms, has motivated cross-pollination of optimization and implementation techniques between the SSM and linear-attention research communities, including shared approaches to efficient parallel computation and, in some proposed architectures, layers that can be understood as interpolating between or combining aspects of both formulations within a single unified framework.

Distillation techniques, discussed in dedicated treatments of model compression, have been applied specifically to convert pretrained Transformer language models into SSM or hybrid architectures without retraining from scratch, initializing the target SSM architecture's parameters based on the source Transformer's learned attention patterns and weights, offering a potentially more compute-efficient path to obtaining a capable SSM-based model than full pretraining from a random initialization.

## Future Research Directions

Closing the remaining quality gap between pure SSM architectures and attention on precise, exact-retrieval-sensitive tasks, whether through improved selective mechanisms, larger or more expressive hidden states, or principled hybrid designs, remains an active area of architecture research, since this gap is currently the primary quality-based argument against replacing attention entirely with SSM layers in applications where such precise retrieval matters.

Further unifying the theoretical understanding of SSMs, linear attention, and related efficient sequence architectures, building on the structured state-space duality result, may reveal additional shared structure that enables cross-architecture transfer of efficiency techniques, training recipes, and theoretical guarantees, reducing the currently substantial duplicated research and engineering effort across what have historically been treated as largely separate architectural lineages.

Extending SSM architectures effectively to multimodal settings (vision, audio, and combined modalities processed jointly) and further scaling studies establishing how SSM and hybrid architectures compare to Transformers as both model size and training compute continue to grow are active, closely watched areas, given that most published comparisons to date, while informative, have not yet spanned the very largest model and data scales at which some architectural trade-offs have historically only become fully apparent.

## Summary & Key Takeaways

State-space models adapt the classical continuous-time linear dynamical systems formalism into a discretized, learnable neural sequence layer, offering linear-time sequence processing and constant per-token inference memory, in contrast to the quadratic time and linearly growing memory cost of standard Transformer self-attention.

The S4 model established that a structured, HiPPO-motivated initialization of the state transition matrix is essential to achieving strong long-range dependency modeling with this formalism, while Mamba's selective state-space formulation, making recurrence parameters input-dependent, substantially closed the expressive gap to attention-based models while preserving efficient recurrent computation via a custom hardware-aware parallel scan.

Pure SSM architectures still show measurable weaknesses relative to attention on tasks demanding precise, exact retrieval from long contexts, motivating hybrid architectures that interleave a small number of attention layers among a majority of SSM layers to capture much of the efficiency benefit of the SSM backbone while mitigating this specific limitation.

SSM and Mamba-based architectures have found particularly strong practical application in domains demanding very long context lengths, including long-document language modeling, genomic sequence modeling, and audio processing, and continue to be an active area of both architectural and theoretical research as a complement and, in some settings, an alternative to standard Transformer attention.

Keywords: state-space model, SSM, S4, Mamba, selective state-space, structured state-space sequence model, HiPPO initialization, linear recurrence, parallel scan, hardware-aware kernel, long-range arena, quadratic attention, linear attention, Mamba-2, structured state-space duality, hybrid attention SSM architecture, needle in a haystack retrieval, key-value cache, discretization, zero-order hold

---

## Appendix: Practical Labs

### Lab 1: Discretizing a Continuous-Time State-Space Model via Zero-Order Hold

import numpy as np

np.random.seed(0)


def zero_order_hold_discretize(A, B, delta):
    """Discretizes a continuous-time state-space system (A, B) with step size delta
    using the zero-order hold method:
        A_bar = exp(delta * A)
        B_bar = (delta * A)^-1 (exp(delta * A) - I) * delta * B
    Falls back to a first-order (Euler) approximation for B_bar when A is singular."""
    from scipy.linalg import expm

    n = A.shape[0]
    A_bar = expm(delta * A)

    try:
        dA_inv = np.linalg.inv(delta * A)
        B_bar = dA_inv @ (A_bar - np.eye(n)) @ (delta * B)
    except np.linalg.LinAlgError:
        # Singular delta*A: fall back to a simple Euler discretization for B
        B_bar = delta * B

    return A_bar, B_bar


def run_discrete_recurrence(A_bar, B_bar, C, inputs):
    """Runs h_t = A_bar h_{t-1} + B_bar x_t, y_t = C h_t over a sequence of scalar inputs."""
    n = A_bar.shape[0]
    h = np.zeros(n)
    outputs = []
    for x_t in inputs:
        h = A_bar @ h + (B_bar.flatten() * x_t)
        y_t = C @ h
        outputs.append(y_t.item() if np.ndim(y_t) > 0 else y_t)
    return np.array(outputs)


def test_discretization_stability_and_output():
    # Stable continuous-time system: eigenvalues of A have negative real parts
    A = np.array([[-0.5, 1.0], [0.0, -0.3]])
    B = np.array([[1.0], [0.5]])
    C = np.array([1.0, 0.0])

    delta = 0.1
    A_bar, B_bar = zero_order_hold_discretize(A, B, delta)

    eigenvalues_discrete = np.linalg.eigvals(A_bar)
    print(f"Discrete A_bar eigenvalue magnitudes: {np.abs(eigenvalues_discrete)}")

    # A stable continuous system discretized with ZOH should yield a discrete
    # system whose eigenvalues lie strictly inside the unit circle (magnitude < 1)
    assert np.all(np.abs(eigenvalues_discrete) < 1.0), (
        "Discretized stable system should have eigenvalues inside the unit circle"
    )

    inputs = np.sin(np.linspace(0, 4 * np.pi, 50)) + np.random.randn(50) * 0.01
    outputs = run_discrete_recurrence(A_bar, B_bar, C, inputs)

    print(f"Output sequence length: {len(outputs)}")
    print(f"Output range: [{outputs.min():.3f}, {outputs.max():.3f}]")

    assert len(outputs) == len(inputs), "Output sequence length should match input length"
    assert np.all(np.isfinite(outputs)), "Discretized recurrence should not diverge for a stable system"

    print("Zero-order hold discretization test passed.")


if __name__ == "__main__":
    test_discretization_stability_and_output()

### Lab 2: Selective (Input-Dependent) vs. Fixed SSM Recurrence Parameters

import numpy as np

np.random.seed(1)


def softplus(x):
    return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0)


class FixedSSM:
    """Non-selective SSM: A, B, C, and step size delta are fixed constants,
    shared across every position in the sequence (S4-style)."""

    def __init__(self, state_dim, delta=0.1, seed=0):
        rng = np.random.RandomState(seed)
        self.A = -np.abs(rng.randn(state_dim)) - 0.1  # stable, negative diagonal dynamics
        self.B = rng.randn(state_dim) * 0.5
        self.C = rng.randn(state_dim) * 0.5
        self.delta = delta

    def forward(self, inputs):
        h = np.zeros_like(self.A)
        outputs = []
        A_bar = np.exp(self.delta * self.A)
        B_bar = self.delta * self.B
        for x_t in inputs:
            h = A_bar * h + B_bar * x_t
            outputs.append(np.dot(self.C, h))
        return np.array(outputs), None


class SelectiveSSM:
    """Selective SSM: delta and B are computed from the current input token,
    allowing the model to modulate how strongly each token affects the state (Mamba-style)."""

    def __init__(self, state_dim, seed=0):
        rng = np.random.RandomState(seed)
        self.A = -np.abs(rng.randn(state_dim)) - 0.1
        self.C = rng.randn(state_dim) * 0.5
        self.w_delta = rng.randn() * 0.5
        self.w_B = rng.randn(state_dim) * 0.5

    def forward(self, inputs):
        h = np.zeros_like(self.A)
        outputs = []
        effective_deltas = []
        for x_t in inputs:
            delta_t = softplus(self.w_delta * x_t) + 1e-3
            B_t = self.w_B * x_t
            A_bar_t = np.exp(delta_t * self.A)
            B_bar_t = delta_t * B_t
            h = A_bar_t * h + B_bar_t * x_t
            outputs.append(np.dot(self.C, h))
            effective_deltas.append(delta_t)
        return np.array(outputs), np.array(effective_deltas)


def test_selective_ssm_modulates_step_size():
    state_dim = 8
    seq_len = 40

    # Construct an input with a mix of near-zero ("unimportant") and
    # large-magnitude ("important") tokens
    inputs = np.zeros(seq_len)
    important_positions = [5, 15, 25, 35]
    inputs[important_positions] = 3.0
    inputs += np.random.randn(seq_len) * 0.05  # small background noise elsewhere

    selective_model = SelectiveSSM(state_dim, seed=2)
    outputs, effective_deltas = selective_model.forward(inputs)

    avg_delta_important = effective_deltas[important_positions].mean()
    avg_delta_other = np.delete(effective_deltas, important_positions).mean()

    print(f"Average effective step size at important positions: {avg_delta_important:.4f}")
    print(f"Average effective step size at other positions: {avg_delta_other:.4f}")

    # The selective mechanism should produce a measurably different (here, larger,
    # since larger input magnitude increases softplus(w_delta * x_t)) step size
    # at high-magnitude "important" tokens compared to near-zero background tokens.
    assert avg_delta_important != avg_delta_other, (
        "Selective SSM should modulate step size differently based on input content"
    )
    assert np.all(np.isfinite(outputs)), "Selective SSM output should remain finite"

    fixed_model = FixedSSM(state_dim, delta=0.1, seed=2)
    fixed_outputs, _ = fixed_model.forward(inputs)
    assert np.all(np.isfinite(fixed_outputs)), "Fixed SSM output should also remain finite"

    print("Selective vs. fixed SSM step-size modulation test passed.")


if __name__ == "__main__":
    test_selective_ssm_modulates_step_size()

### Lab 3: Sequential Recurrence vs. Parallel Scan Equivalence

import numpy as np

np.random.seed(2)


def sequential_scan(A_bar_seq, B_bar_x_seq):
    """Computes h_t = A_bar_t * h_{t-1} + B_bar_x_t sequentially, one step at a time.
    A_bar_seq and B_bar_x_seq are arrays of shape (seq_len, state_dim)."""
    seq_len, state_dim = A_bar_seq.shape
    h = np.zeros(state_dim)
    states = []
    for t in range(seq_len):
        h = A_bar_seq[t] * h + B_bar_x_seq[t]
        states.append(h.copy())
    return np.stack(states)


def associative_combine(left, right):
    """Combines two (A, Bx) pairs under the associative operator used in a parallel scan
    for linear recurrences: (A1, Bx1) . (A2, Bx2) = (A2 * A1, A2 * Bx1 + Bx2)."""
    A1, Bx1 = left
    A2, Bx2 = right
    return A2 * A1, A2 * Bx1 + Bx2


def parallel_scan(A_bar_seq, B_bar_x_seq):
    """A simple (non-work-optimal, but correctness-equivalent) parallel scan implementation
    using a sequential left-to-right reduction with the associative combine operator,
    which produces identical results to the sequential scan but illustrates the
    associative structure that real parallel scan implementations exploit for
    logarithmic-depth (rather than linear-depth) computation on parallel hardware."""
    seq_len, state_dim = A_bar_seq.shape
    combined_A = A_bar_seq[0].copy()
    combined_Bx = B_bar_x_seq[0].copy()
    states = [combined_Bx.copy()]

    for t in range(1, seq_len):
        combined_A, combined_Bx = associative_combine(
            (combined_A, combined_Bx), (A_bar_seq[t], B_bar_x_seq[t])
        )
        states.append(combined_Bx.copy())

    return np.stack(states)


def test_parallel_scan_matches_sequential_scan():
    seq_len, state_dim = 30, 6

    # Stable, input-dependent (selective-style) A_bar values in (0, 1)
    A_bar_seq = np.random.uniform(0.5, 0.95, size=(seq_len, state_dim))
    B_bar_x_seq = np.random.randn(seq_len, state_dim) * 0.3

    seq_states = sequential_scan(A_bar_seq, B_bar_x_seq)
    par_states = parallel_scan(A_bar_seq, B_bar_x_seq)

    max_abs_diff = np.max(np.abs(seq_states - par_states))
    print(f"Sequence length: {seq_len}, state dim: {state_dim}")
    print(f"Max absolute difference between sequential and parallel scan: {max_abs_diff:.2e}")

    assert max_abs_diff < 1e-8, (
        "Parallel scan (via associative combine) must produce numerically "
        "identical results to the sequential scan"
    )

    print("Sequential vs. parallel scan equivalence test passed.")


if __name__ == "__main__":
    test_parallel_scan_matches_sequential_scan()

### Lab 4: Constant-Memory Autoregressive Inference vs. Growing KV-Cache Simulation

import numpy as np

np.random.seed(3)


def ssm_autoregressive_memory_usage(state_dim, n_generated_tokens):
    """An SSM maintains a single fixed-size hidden state regardless of how many
    tokens have been generated: memory usage per step is constant."""
    memory_per_step = []
    for step in range(1, n_generated_tokens + 1):
        # The hidden state size never changes as generation proceeds
        memory_per_step.append(state_dim)
    return np.array(memory_per_step)


def attention_autoregressive_memory_usage(head_dim, n_generated_tokens):
    """A Transformer with standard attention must retain a key and value vector
    for every previously generated token in its KV cache, so memory grows
    linearly with the number of tokens generated so far."""
    memory_per_step = []
    for step in range(1, n_generated_tokens + 1):
        # KV cache holds one key + one value vector (each of size head_dim) per past token
        memory_per_step.append(step * 2 * head_dim)
    return np.array(memory_per_step)


def test_constant_vs_growing_memory():
    state_dim = 64
    head_dim = 64
    n_tokens = 1000

    ssm_memory = ssm_autoregressive_memory_usage(state_dim, n_tokens)
    attn_memory = attention_autoregressive_memory_usage(head_dim, n_tokens)

    print(f"SSM memory usage: constant at {ssm_memory[0]} units for all {n_tokens} steps")
    print(f"Attention KV-cache memory at step 1: {attn_memory[0]} units")
    print(f"Attention KV-cache memory at step {n_tokens}: {attn_memory[-1]} units")

    # The SSM's memory footprint must not change at all across generation steps
    assert np.all(ssm_memory == ssm_memory[0]), "SSM memory usage should be exactly constant"

    # The attention KV-cache must grow strictly monotonically as more tokens are generated
    assert np.all(np.diff(attn_memory) > 0), "Attention KV-cache memory should grow with each step"

    # At long generation lengths, the attention KV-cache should vastly exceed
    # the SSM's fixed memory footprint
    ratio_at_end = attn_memory[-1] / ssm_memory[-1]
    print(f"Attention/SSM memory ratio at final step: {ratio_at_end:.1f}x")
    assert ratio_at_end > 10.0, (
        "At long context lengths, attention KV-cache memory should dwarf constant SSM memory"
    )

    print("Constant-memory SSM vs. growing-KV-cache attention test passed.")


if __name__ == "__main__":
    test_constant_vs_growing_memory()