File size: 2,401 Bytes
1cd8a52 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | ## Code for "HiPPO Zoo: Explicit Memory Mechanisms for Interpretable State Space Models" [ICML 2026]
[](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}
}
```
|