🤖 machine learning

Flash PD-SSM: Memory-Optimized Structured Sparse State-Space Models

The paper introduces Flash PD-SSM, a memory-optimized structured sparse state-space model that dynamically selects from a set of trainable sparse matrices to achieve the high expressivity of unstructured models while maintaining the computational efficiency required for large-scale training and state-of-the-art performance on long sequences.

Original authors: Aleksandar Terzić, Francesco Carzaniga, Nicolas Menet, Yannick Biehl, Michael Hersche, Thomas Hofmann, Abbas Rahimi

Published 2026-05-20

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Aleksandar Terzić, Francesco Carzaniga, Nicolas Menet, Yannick Biehl, Michael Hersche, Thomas Hofmann, Abbas Rahimi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to build a super-smart robot that can read long stories, remember details, and predict what happens next. To do this, the robot needs a "brain" that can hold onto information as it reads.

For a long time, the best brains for this job were like Transformers (the technology behind many current AI chatbots). They are incredibly smart but very heavy and slow, like a luxury limousine that guzzles gas and needs a huge garage.

Then, engineers invented State-Space Models (SSMs). Think of these as electric scooters. They are much faster and use way less energy (memory), making them perfect for long trips. However, early scooters had a problem: they were too simple. They could handle straight roads, but if the path got twisty or required complex logic (like remembering a specific rule for a specific character), they would get lost.

The paper introduces a new invention called FLASH PD-SSM. It's like taking that electric scooter and upgrading it with a smart, modular navigation system that makes it as good at handling complex twists as the heavy limousine, but without losing its speed or fuel efficiency.

Here is how they did it, broken down into simple concepts:

1. The Problem: The "Too-Heavy" Map

Previous versions of these smart scooters (like the model called PD-SSM) tried to be super smart by carrying a massive, detailed map for every single step of the journey.

The Issue: Carrying this huge map for a long journey took up so much memory that the scooter would run out of battery before it got far. It was too heavy to be practical for real-world use.
The Result: Other models (like Mamba) chose to carry a tiny, simple map. They were fast and light, but they couldn't solve complex puzzles or remember intricate rules.

2. The Solution: The "Magic Switch"

The authors of this paper, FLASH PD-SSM, came up with a clever trick. Instead of carrying a massive, detailed map for every step, they built a toolbox of pre-made, specialized maps.

How it works: At every step of the journey, the robot looks at the current situation and flips a switch to pick the one perfect map from the toolbox that fits that moment.
The Analogy: Imagine you are driving. Instead of drawing a new, detailed map of the whole city every time you turn a corner (which takes forever and uses a lot of paper), you have a set of 100 pre-drawn "local maps." You just pick the one you need for the next block.
The Benefit: This switch is incredibly fast to flip and takes up almost no space. This allows the robot to have the complexity of the heavy limousine (it can solve hard logic puzzles) but keeps the speed and lightness of the scooter.

3. What They Proved

The team tested this new "smart scooter" in three main ways:

The Logic Test (FSA Emulation): They gave the robot a series of logic puzzles (like counting parity or navigating a maze). The FLASH PD-SSM solved almost all of them (96% accuracy), while the simpler models (like Mamba2) struggled significantly. It proved the robot could actually "think" through complex rules.
The Long Memory Test (Time Series): They asked the robot to analyze data streams that were incredibly long (over 17,000 steps). FLASH PD-SSM was the most accurate at predicting what would happen next, beating all other competitors.
The Language Test: They put this new brain into a language model (a robot that writes text).
- State Tracking: When asked to track objects moving between boxes in a story (e.g., "The red ball moved from Box A to Box B"), the new model was much better at remembering the final location than previous models.
- Writing: When trained to write like a human, a hybrid version of this model (mixing the new brain with existing tech) wrote better and faster than models using only the older, simpler brains.

4. The Bottom Line

The paper claims that FLASH PD-SSM solves the biggest trade-off in AI right now: Speed vs. Intelligence.

Before: You had to choose between being fast but "dumb" (simple models) or being "smart" but slow and expensive (complex models).
Now: FLASH PD-SSM is like a high-speed train that can also climb mountains. It runs as fast as the current top models (Mamba2) but can handle much more complex tasks, all while using less memory.

In short, they found a way to make AI models that are lighter, faster, and smarter all at the same time, without needing to build a bigger, more expensive computer to run them.

Technical Summary: FLASH PD-SSM

Problem Statement

State-space models (SSMs) offer a scalable alternative to Transformers, providing linear compute scaling and constant memory consumption during inference. However, they face a fundamental trade-off between efficiency and expressivity, dictated by the structure of their transition matrices.

Unstructured matrices offer maximal expressivity (capable of emulating any Finite State Automaton, or FSA) but incur prohibitive compute and memory costs.
Structured matrices (e.g., diagonal, scaled identity) are highly efficient but suffer from limited expressivity, theoretically incapable of emulating FSAs with non-solvable transformation semigroups.
Structured Sparse Matrices (PD-SSM): A recent advancement (PD-SSM) utilizes structured sparse matrices to achieve FSA-level expressivity while enabling efficient parallel computation. However, PD-SSM suffers from a critical bottleneck: generating the input-dependent sparse matrices incurs an $O(LN^2)$ memory cost (where $L$ is sequence length and $N$ is hidden state size). This overhead makes PD-SSM significantly slower (40 $\times$ ) and more memory-intensive than efficient competitors like Mamba2, rendering it unsuitable for large-scale deployment.

Methodology: FLASH PD-SSM

The authors propose FLASH PD-SSM, a memory-optimized structured sparse SSM that resolves the generation bottleneck of PD-SSM while preserving its theoretical expressivity guarantees.

1. Novel Selection Mechanism (Hard Selection)

The core innovation replaces the soft-selection mechanism of PD-SSM with a hard selection mechanism:

Original PD-SSM: Generates a convex combination of dense dictionary matrices and applies a column-wise hardmax to sparsify the result. This requires materializing dense $N \times N$ matrices for every time step, leading to $O(LN^2)$ memory.
FLASH PD-SSM: Pre-computes a dictionary of $K$ sparse matrices. At each time step, the input $u_t$ generates $K$ scores via a linear transformation. The model selects the single matrix with the highest score (argmax) and uses it directly.
Complexity Reduction: Since the selected matrix is already sparse (a one-hot column selector), the memory cost drops from $O(LN^2)$ to $O(LN)$. The compute cost for matrix generation reduces from $O(LN^2K)$ to $O(LNK)$.

2. Gradient Propagation

Because the selection involves discrete operations (argmax/hardmax), the model employs Straight-Through Estimators (STE) with slope annealing for backpropagation:

In the forward pass, hard selection (argmax) is used.
In the backward pass, the hard selection is approximated by a tempered softmax to allow gradient flow through the selection weights and the dictionary matrices.
Proposition 2 in the paper provides the surrogate gradient expressions, which become exact as the temperature parameter $\tau \to 0$ .

3. Hardware-Aware Implementation

The authors implement a custom three-phase CUDA kernel (wrapped in Triton) to exploit GPU parallelism:

Phase A: Computes independent sequential recurrences within chunks of the sequence.
Phase B: Aggregates inter-chunk dependencies (prefix scan) to compute carry states.
Phase C: Applies the carry corrections to the chunk states.
This design minimizes memory traffic and maximizes bandwidth utilization, achieving near-peak GPU memory bandwidth.

4. Architecture Integration

FLASH PD-SSM is embedded in a standard pre-norm architecture with multi-head SSM layers. The multi-head design allows for the amortization of parameter costs for the dense dictionary matrices. The authors demonstrate that scaling the dictionary size $K$ as $O(\sqrt{D})$ (where $D$ is embedding dimension) keeps the parameter count manageable without "blowing up" the model size.

Key Contributions

Algorithmic Improvement: Reduces the memory complexity of structured sparse SSM generation from $O(LN^2)$ to $O(LN)$ via a discrete selection mechanism, enabling training at scale.
Theoretical Expressivity: Proves (Proposition 1) that FLASH PD-SSM retains the theoretical guarantee of unstructured SSMs: it can exactly represent any deterministic FSA with $N$ states using a single layer of size $N$ .
Hardware Optimization: Delivers a custom CUDA implementation that achieves state-of-the-art throughput and memory efficiency.
Hybrid Model Integration: Demonstrates that FLASH PD-SSM serves as an effective drop-in replacement in hybrid LLMs, improving state-tracking capabilities without sacrificing throughput.

Experimental Results

Synthetic State-Tracking & Expressivity

FSA Emulation: On synthetic tasks (parity, modular counting, cycle navigation), FLASH PD-SSM achieves 96.3% average accuracy, significantly outperforming diagonal SSMs (Mamba2: 58.9%, Mamba3: 90.7%) and DPLR models (Gated DeltaProduct: 80.7%). This confirms the practical realization of its theoretical expressivity.
Mechanistic Capabilities (MAD Benchmark): FLASH PD-SSM scores 69.2%, comparable to Mamba (69.3%), indicating strong in-context recall and memorization capabilities.

Runtime and Memory Efficiency

Throughput: At a sequence length of 5120, FLASH PD-SSM is 18% faster than DeltaNet and 4% faster than Mamba2.
Memory: It reduces peak memory consumption by >12% compared to both DeltaNet and Mamba2 at sequence lengths above 4000.
Scalability: The model maintains high throughput across varying sequence lengths, addressing the bottleneck that previously hindered PD-SSM.

Natural Language State-Tracking

Evaluated on the Box dataset (extended to 26 state transitions) using hybrid models (Granite4, Nemotron-H).
FLASH PD-SSM layers inserted into these models improved out-of-distribution (OOD) accuracy by up to 10% over Mamba2 and 8% over DeltaNet, while maintaining comparable throughput.

Language Modeling

DELTA-PD: A hybrid model combining FLASH PD-SSM and DeltaNet (1:3 ratio) was pre-trained on FineWeb100B.
Performance: DELTA-PD surpassed pure DeltaNet in accuracy (42.3% vs 41.3% average) and throughput (253k tok/s vs 168k tok/s, a 1.5 $\times$ improvement).
It also outperformed pure FLASH PD-SSM, suggesting that hybridization balances the strengths of different SSM types.

Time-Series Modeling

Achieved State-of-the-Art (SoTA) accuracy on the UEA multivariate time-series classification task (sequences >17,000).
Reached near-SoTA Mean Squared Error (MSE) on the PPG-DaLiA forecasting task.

Significance and Claims

The paper claims that FLASH PD-SSM successfully breaks the efficiency-expressivity trade-off that has limited SSMs. By reducing the memory overhead of structured sparse matrices, it enables the use of highly expressive models (capable of FSA emulation) at the scale required for modern language models.

The authors position FLASH PD-SSM not just as a theoretical improvement, but as a practical, drop-in replacement for existing SSM layers in frontier hybrid models. It offers a unique position in the performance landscape: it is the only SSM identified in the "upper right quadrant" of Figure 1, meaning it is simultaneously highly expressive and fast/memory-efficient.

The paper acknowledges limitations regarding the use of straight-through estimators, noting that the theoretical explanation for their effectiveness in highly discrete sequential methods remains an open area for future research. However, the empirical results across synthetic, mechanistic, and natural language tasks validate the approach as a robust solution for scalable, expressive sequence modeling.