## Code for "HiPPO Zoo: Explicit Memory Mechanisms for Interpretable State Space Models" [ICML 2026]
[![arXiv](https://img.shields.io/badge/arXiv-2602.21340-b31b1b.svg)](https://arxiv.org/pdf/2602.21340.pdf) 

- Authors: [Jack Goffinet](https://jackgoffinet.github.io),  [Casey Hanks](https://scholar.google.com/citations?hl=en&user=1qPlymwAAAAJ), [David Carlson](https://carlson.pratt.duke.edu/)

<div align="center">

<img src="teaser.jpg" width="600">

</div>


### 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:

```bash
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
```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}
}
```