neural-ssm

Robust neural state-space models in PyTorch.neural-ssm combines recurrent state-space layers, bounded nonlinearities, and optional context-aware routing to build sequence models with an explicit zero-state input/output L2-gain bound.

It is useful when a model must be both expressive and well behaved: nonlinear system identification, disturbance-to-control maps, streaming sequence models, and robustness experiments.

At a glance

You need	Use
A standard stable LTI baseline	DeepSSM(..., param="lru")
A deep model with a prescribed L2 certificate	A certified param, a bounded ff, and gamma=...
Selective, input-dependent dynamics	param="tv" or "tvc"
Safe conditioning on context	ContextualDeepSSM with input, gate, mixer, and/or select ports
System-ID comparisons and visual reports	Test_files/run_benchmarks.py
CUDA throughput	parallel scan/FFT modes, CUDA graphs, and torch.compile

Model flow

flowchart LR
U["input sequence u"] --> E["encoder"] --> B["stack of SSL blocks"] --> D["decoder"] --> Y["output sequence y"]
B --- R["recurrent SSM\n+plus bounded feed-forward branch"]
C["optional context z"] --> P["input / gate / mixer / select ports"]
P --> B

Loading

An SSL block has separate residual temporal and channel-mixing branches. With a certified recurrent cell, a bounded feed-forward layer, and a prescribed gamma, the stack reports a conservative zero-state L2 gain bound.

Install

The core package requires Python 3.9+ and PyTorch 2.2+.

git clone https://github.com/DecodEPFL/SSM.git
cd SSM
python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -e .

For the tutorials, plotting utilities, and nonlinear system-identification benchmark harness, install the experiment extra:

python -m pip install -e ".[experiments]"

For development and the automated tests:

python -m pip install -e ".[dev,experiments]"
python -m pytest

For CUDA, install the PyTorch build appropriate for the target CUDA runtime before installing this project. The package does not install a CUDA runtime on its own.

Quick start

DeepSSM consumes a tensor shaped (batch, time, features) and returns an output sequence plus one final state per SSM block.

importtorchfromneural_ssmimportDeepSSMmodel=DeepSSM(
d_input=3,
d_output=2,
d_model=32,
d_state=32,
n_layers=2,
param="tv", # selective, L2-bounded recurrent cellff="MBLIP", # bounded feed-forward branchgamma=1.5, # prescribed zero-state L2-gain bound
)
u=torch.randn(8, 256, 3)
y, state=model(u, mode="scan")
print(y.shape) # torch.Size([8, 256, 2])print(model.certified_gain_bound())

For streaming inference, retain the returned state and disable the reset:

y_1, state=model(u[:, :128], mode="scan")
y_2, state=model(u[:, 128:], state=state, mode="scan", reset_state=False)

Use detach_state=False only when intentionally backpropagating through multiple calls (cross-call BPTT).

Choose a recurrent core

param	Core	Certificate	Best execution mode
lru	complex diagonal stable LTI recurrence	stable, but no global DeepSSM bound	scan or conv
l2ru	free L2-bounded LTI recurrence	yes	scan or loop
zak	constrained L2-bounded LTI recurrence	yes	scan or loop
l2n	2×2-block L2-bounded LTI recurrence	yes	scan or conv
l2nt	dense L2-bounded LTI recurrence	yes	scan or loop
tv	selective diagonal SSM	yes	scan
tvc	selective LTI SSM	yes	scan

l2n uses 2×2 state blocks, so d_state must be even. The benchmark harness selects the fastest supported mode by default: convolution for lru/l2n and parallel scan for the other SSMs.

Certificates in practice

Set gamma to request a global certificate. This requires a certified recurrent core (l2ru, zak, l2n, l2nt, tv, or tvc) and a feed-forward layer with a declared global bound: LGLU2, BLGLU2, MBLIP, or TLIP. Keep learn_x0=False: a learned nonzero initial state needs a separate storage-energy term and is not covered by the pure induced-gain statement.

The bound applies to zero-state sequence maps. It is a conservative guarantee, not a promise that every trained model will use the full gain budget.

Context-aware DeepSSMs

ContextualDeepSSM wraps a normal DeepSSM and lets a second sequence shape the model without losing the core's certificate.

fromneural_ssmimportContextualDeepSSMcontroller=ContextualDeepSSM(
d_input=3,
d_context=2,
d_output=2,
context_modes=("input", "gate", "mixer", "select"),
context_filter="difference",
d_features=16,
mixer_bound=0.8,
d_model=32,
d_state=32,
n_layers=2,
param="tv",
ff="MBLIP",
gamma=1.5,
)
disturbance=torch.randn(4, 200, 3)
context=torch.randn(4, 200, 2)
correction, state=controller(disturbance, context, mode="scan")

The four ports are complementary:

Port	Effect	When it fits
input	Filters then concatenates context to the SSM input	exogenous references or finite-horizon context
gate	Uses context gates in [0, 1] to attenuate residual branches	context-dependent modulation
mixer	Applies a uniformly bounded context-dependent output matrix	endogenous or in-loop context
select	Conditions selective-cell parameters directly	tv/tvc dynamics that adapt to context

See the contextual tutorial for the filter choices and the corresponding gain diagnostics.

Additional recurrent model

REN is a robust acyclic recurrent-equilibrium network for system identification. Public entry points include DeepSSM, SSMConfig, ContextualDeepSSM, REN, LRU, L2RU, lruz, and the neural_ssm.ssm / neural_ssm.layers namespaces.

Speed on CUDA

The library offers three compatible acceleration layers:

1. Parallel scan and FFT convolution: choose mode="scan" or mode="conv" for sequence-parallel execution where supported.
2. CUDA graphs for selective scans: set use_cuda_graph=True in SSMConfig for fixed-shape tv/tvc workloads. The benchmark harness enables this by default.
3. torch.compile: compile the model after construction and before its first call. The main DeepSSM, contextual, and complex LTI paths are covered by compile regression tests.

device=torch.device("cuda")
model=model.to(device)
compiled_model=torch.compile(model, mode="reduce-overhead")
output, state=compiled_model(u.to(device), mode="scan")

Compilation has a warm-up cost and specializes to observed execution patterns; keep batch and sequence shapes stable when throughput matters. REN.forward is intentionally eager, and the benchmark harness leaves LSTM/GRU/REN uncompiled because their cuDNN/eager paths are the better default.

Run the included experiments

Tutorials

python Test_files/Tutorial_DeepSSM.py
python Test_files/Tutorial_ContextualSSM.py

Nonlinear system-identification benchmark harness

The harness trains complete native trajectories, evaluates the benchmark's official validation split, and writes figures, GIFs, Markdown, and JSON reports.

# Discover datasets and model choices
python Test_files/run_benchmarks.py --list
# A small CUDA run with automatic SSM-mode selection
python Test_files/run_benchmarks.py \
--benchmarks Cascaded_Tanks \
--models tv tvc lru lstm \
--device cuda --epochs 200 --compile --compile-mode reduce-overhead
# Launch the desktop benchmark form (requires Tk support)
python Test_files/benchmark_ui.py

For a focused L2RU-versus-ZAK study with published-style figures:

python scripts/compare_l2ru_vs_zak.py --datasets Cascaded_Tanks --epochs 8

Repository map

src/neural_ssm/
├── ssm/ DeepSSM, recurrent cells, scan utilities, and context models
├── static_layers/ bounded and conventional feed-forward layers
└── rens/ robust acyclic REN
tests/ certificate, context, and torch.compile regression tests
Test_files/ tutorials, benchmark runner, visual UI, and research scripts
scripts/ focused reproducible experiment drivers
docs/ figures and deployment notes

Citation

If you use this repository in research, please cite:

Free Parametrization of L2-bounded State Space Models

https://arxiv.org/abs/2503.23818