David G. Clark1,*, Blake Bordelon2,*, Jacob A. Zavatone-Veth3,4,*, Cengiz Pehlevan1,4,5
1 Kempner Institute for the Study of Natural and Artificial Intelligence, 2 Center of Mathematical Sciences and Applications, 3 Society of Fellows, 4 Center for Brain Science, 5 John A. Paulson School of Engineering and Applied Sciences, Harvard University, Cambridge, MA, USA
* DGC, BB, and JAZ-V contributed equally to this work.
Paper: bioRxiv 2026.03.02.708943
Data: Zenodo 10.5281/zenodo.18831815
dmft_solver.py # DMFT equations: solve_dmft, dmft_action, sample_dmft_trajectories
rnn_util.py # Batched RNN forward pass, SGLD training step, energy computation
plot_utils.py # Shared plotting: eigenvalue sizing, ridgeline, participation ratio
monkey/ # Motor cortex RNN experiments
analysis.ipynb # Paper figures notebook
helpers.py # Data loading, eigenvalue analysis, similarity, selectivity
similarity.py # SVCCA and CKA metrics
style.py # Matplotlib styling
train.py # SGLD training script
train.sh # SLURM submission script
sine_wave/ # Sine wave DMFT experiments
analysis.ipynb # Visualization notebook (SGLD + DMFT theory)
train.py # Batched SGLD training + DMFT solver
train.sh # SLURM submission script
linear/ # Linear RNN theory
analysis.ipynb # Linear RNN analysis notebook
theory.py # Analytical theory for linear networks
sussillo_data/ # Neural and behavioral data from Sussillo et al. (2015)
README.md # Data documentation (variables, MATLAB struct fields, provenance)
load_sussillo_data.py # Generator script: raw MATLAB → sussillo_data.npz
Note on terminology: The paper refers to "Langevin gradient flow" while the codebase uses "SGLD" (stochastic gradient Langevin dynamics). We use "flow" in the paper since there are no minibatches and hence no stochasticity in the traditional sense, but the code inherits the SGLD naming.
Download all archives from Zenodo.
Trained model outputs. Extract into a data/ directory at the repository root:
mkdir -p data
tar xf monkey_results.tar -C data/
tar xf monkey_weights_and_cache.tar -C data/
tar xf sine_wave_and_linear.tar -C data/This gives the directory layout expected by the analysis notebooks:
data/
monkey/
results_000.npz ... results_009.npz # Trained RNN outputs
weights_000.pt ... weights_009.pt # Final weight matrices
weights_*_ckpt400000.pt # Checkpoint weight matrices
cache/ # Precomputed analysis caches
sine_wave/
results_dmft_*.npz # DMFT + SGLD results
cache/ # DMFT trajectory cache
linear/
cache/ # Participation ratio cache
Neural and behavioral data. The sussillo_data.tar archive contains both the processed file (sussillo_data.npz) and the raw MATLAB source files. Extract and place the .npz at the repository root (where the notebooks expect it):
tar xf sussillo_data.tar
cp sussillo_data/sussillo_data.npz .The raw MATLAB files will be at sussillo_data/churchland_sussillo_data/. The generator script sussillo_data/load_sussillo_data.py can regenerate the .npz from these files if needed.
If you only need the neural and behavioral data (without trained models), download just sussillo_data.tar from Zenodo.
See sussillo_data/README.md for variable descriptions and data provenance.
Open and run the Jupyter notebooks:
monkey/analysis.ipynb— Motor cortex RNN figuressine_wave/analysis.ipynb— Sine wave / DMFT figureslinear/analysis.ipynb— Linear RNN theory figures
To retrain RNNs on a SLURM cluster:
sbatch monkey/train.sh # 10 GPU jobs
sbatch sine_wave/train.sh # 20 GPU jobs- Python 3.10+
- PyTorch 2.0+
- NumPy, SciPy, Matplotlib
- Jupyter
All data needed to reproduce the figures are available on Zenodo:
- Motor cortex RNN outputs (
monkey_results.tar,monkey_weights_and_cache.tar) — activations, preactivations, eigenvalue spectra, and loss traces for trained motor cortex RNNs, plus full weight matrices and precomputed analysis caches (similarity metrics, selectivity, etc.) - Sine wave and linear theory outputs (
sine_wave_and_linear.tar) — DMFT solutions, SGLD-trained RNN outputs, eigenvalue spectra, and cached DMFT trajectories for the sine wave task, plus participation ratio data for the linear theory - Neural and behavioral data (
sussillo_data.tar) — trial-averaged firing rates from M1/PMd, EMG targets, and hand trajectories
The neural recordings are from Churchland et al. (2012), as used in Sussillo et al. (2015). The RNN training data (inputs and EMG targets) are from Sussillo et al. (2015).
We are indebted to Mark Churchland for kindly sharing data from Sussillo et al. (2015) and for allowing us to publicly share the full dataset, and for helpful discussions. We thank Juan Carlos Fernández del Castillo for helpful comments on a previous version of this manuscript.
In the initial phase of this work, B.B. was supported by a Google Ph.D. Fellowship, and J.A.Z.-V. and C.P. were supported by NSF Award DMS-2134157, NSF CAREER Award IIS-2239780, and a Sloan Research Fellowship. At present, D.G.C. is supported by the Kempner Institute for the Study of Natural and Artificial Intelligence at Harvard University. B.B. is supported by the Center of Mathematical Sciences and Applications at Harvard University. J.A.Z.-V. is supported by the Office of the Director of the National Institutes of Health under Award Number DP5OD037354, and by a Junior Fellowship from the Harvard Society of Fellows. C.P. is presently supported by NSF CAREER Award IIS-2239780, DARPA grants DIAL-FP-038 and AIQ-HR00112520041, and The William F. Milton Fund from Harvard University. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence.
This project is licensed under the MIT License. See LICENSE for details.