source_code/README.md

Code for "HiPPO Zoo: Explicit Memory Mechanisms for Interpretable State Space Models" [ICML 2026]

• Authors: Jack Goffinet, Casey Hanks, David Carlson

Installation

Use a conda enviornment with JAX installed to run most of the code. The exception is the LSTM, Transformer, and S4D model runs, which we recommend running from a separate enviornment with PyTorch installed. Install the package in editable mode together with the optional development requirements for testing:

pip install -e .[dev]

Repository layout

src/mpm/
    __init__.py           # Re-exports the public API
    integration.py        # NumPy/SciPy integrator with an optional trajectory
    polynomials.py        # Legendre and Fourier polynomial helpers
    synthetic.py          # Synthetic datasets
    utils.py              # Shared helper functions
    models.py             # HiPPO Zoo models (JAX)
    lstm.py               # LSTM model (PyTorch)
    transformer.py        # Transformer model (PyTorch)
    s4d.py                # S4D model (PyTorch)
tests/                    # Small tests
scripts/                  # Experiment and figure scripts used in the paper
    multiscale_hippo.py             # Multiscale HiPPO figures
    forecasting_hippo.py            # Forecasting HiPPO figures
    associative_memory_hippo.py     # Associative Memory HiPPO figures
    salience_hippo.py               # Salience HiPPO figures
    salience_hippo_schematic.py     # Salience HiPPO schematic figure
    volterra_hippo.py               # Volterra HiPPO figure
    run_arsc_experiment.py          # JAX benchmark of HiPPOs on AR and SC tasks
    run_arsc_experiment_torch.py    # PyTorch benchmark of baselines on AR and SC tasks

Paper BibTex

@inproceedings{Goffinet2026HiPPOZoo,
  author       = {Goffinet, Jack and Hanks, Casey and Carlson, David E.},
  title        = {HiPPO Zoo: Explicit Memory Mechanisms for Interpretable State Space Models},
  booktitle    = {International Conference on Machine Learning},
  year         = {2026},
  organization = {PMLR}
}