Skip to content

davidclark1/RNN-Learning-Theory

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

2 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Structure, disorder, and dynamics in task-trained recurrent neural circuits

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

Repository structure

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.

Reproducing results

1. Download data

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.

2. Run analysis notebooks

Open and run the Jupyter notebooks:

  • monkey/analysis.ipynb — Motor cortex RNN figures
  • sine_wave/analysis.ipynb — Sine wave / DMFT figures
  • linear/analysis.ipynb — Linear RNN theory figures

3. Retraining from scratch (optional)

To retrain RNNs on a SLURM cluster:

sbatch monkey/train.sh      # 10 GPU jobs
sbatch sine_wave/train.sh   # 20 GPU jobs

Requirements

  • Python 3.10+
  • PyTorch 2.0+
  • NumPy, SciPy, Matplotlib
  • Jupyter

Data availability

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).

Acknowledgements

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.

License

This project is licensed under the MIT License. See LICENSE for details.

About

No description, website, or topics provided.

Resources

License

Stars

3 stars

Watchers

0 watching

Forks

Packages

 
 
 

Contributors