Grid-SSM: A Native & Lightweight PyTorch State Space Model

Project Overview

Grid-SSM is an experimental sequence modeling architecture designed to explore an alternative to State Space Models (SSMs) using pure PyTorch.

Some recent SSMs such as Mamba rely on hidden state expansions (4D tensors), which increase VRAM consumption in standard PyTorch unless optimized with custom CUDA/Triton kernels.

Grid-SSM proposes a different approach. By using a 3D spatial state with dynamic gates and a learned residual mechanism, Grid-SSM aims to offer comparable expressive power while remaining PyTorch native, memory-efficient, and parallelizable.

For an in-depth breakdown of how the architecture works, please read our Architectural Documentation.

Latent Space Visualization

By reducing the native spatial state into geometric coordinates using Principal Component Analysis (PCA), we can visualize the model's "thought trajectory" in real-time as it generates text.

---

Experimental Results

To evaluate Grid-SSM, we conducted a character-level language modeling test comparing it to a Transformer, Mamba2, and Gated DeltaNet. All models were adjusted to have around ~60,000 parameters for a fair comparison.

Author's Note: Due to hardware constraints (lack of server-grade GPUs), scaling tests on larger parameter counts were not possible. These results serve as a small-scale proof of concept.

Model	Parameters	Training Epochs	Final Loss	Final Perplexity	Math Complexity (Time)	Space Complexity
Grid-SSM	~60K	50	1.2018	3.3261	$O(L \cdot D)$ *	$O(L \cdot D)$
Mamba2	~60K	20	1.1954	3.3048	$O(L \cdot D \cdot N)$ *	$O(L \cdot D \cdot N)$
Transformer	~60K	20	1.9756	7.2109	$O(L^2 \cdot D)$	$O(L^2 \cdot D)$
Gated DeltaNet	~60K	-	N/A	N/A	$O(L \cdot D^2)$	$O(D^2)$

* Note: Both Grid-SSM and Mamba2 can be parallelized in $O(\log L)$ time using torch.cumsum or custom parallel scans.

Note on Gated DeltaNet: Excluded from the final comparative graphs because its reference implementation relies on recurrent matrix-vector multiplications. This sequential dependency makes training impractically slow in standard PyTorch without specialized kernels.

Convergence Plot

As shown below, Grid-SSM approaches Mamba2's final performance with a reduced memory footprint.

Architecture Characteristics

Strengths

• Reduced Spatial Complexity: Grid-SSM does not use the $N$ (state) multiplier. It manages sequence history within the native channel dimension.
• PyTorch Native: The architecture uses standard PyTorch operations without requiring custom C++ or Triton kernels.
• Parallelizable: The linear recurrence allows training using PyTorch's torch.cumsum parallel scan.
• Modeling Capacity: It reaches a Perplexity/Loss plateau close to Mamba2 for an equivalent parameter count.

Cons

• Slower Initial Convergence: Mamba2 benefits from an explicit exponential decay initialization that acts as an inductive bias. Grid-SSM learns this gating behavior dynamically, requiring more epochs to converge in our tests (50 epochs vs 20).

Usage

To reproduce the results using uv:

1. Install dependencies:

uv sync

2. Run the training benchmark:

uv run train.py

License

This project is licensed under the Apache 2.0 License.