Training and analysis of Multi-Plastic Networks (MPNs) on a battery of cognitive tasks. A single recurrent network with Hebbian-like synaptic plasticity learns many tasks at once; the structure of its plastic weights is then analyzed to understand how task-specific computation is organized.
The core model is DeepMultiPlasticNet (core/mpn.py): a recurrent
network whose effective weights are modulated by a fast plasticity matrix M:
W_eff(t) = W + W ⊙ M(t) (multiplicative) or W + M(t) (additive)
M evolves by a Hebbian rule with learning rate η and decay λ (each scalar,
pre/post-vector, or full matrix). In every training script η is learned and λ
is fixed at λ = 1 − Δt/m_time_scale, with Δt = 40 ms. The network has three
weight matrices: W_initial_linear (input projection), mp_layer1.W
(recurrent plastic weights) and W_output (readout).
| Experiment | hidden | batch | L2 | m_time_scale (λ) |
|---|---|---|---|---|
multi-task (multiple_task/) |
300 | 128 | 1e-4 (cohorts from 1e-2 to 1e-5) | 4000 ms (0.99) |
single-task (one_task/) |
200 | 128 | 1e-4 | 400 ms (0.9) |
two-task (two_task/) |
200 | 128 | 1e-3 | 4000 ms (0.99) |
pretraining (pretrain/) |
200 | 128 | 1e-3 | 400 ms (0.9); historical cohort 4000 ms (0.99) |
Hyperparameters are set at the top of each training script and encoded in the run name (see Outputs). SCHEME.md explains the timing convention, where each model's time constant lives and the figure colors.
The multi-task ruleset everything holds 15 tasks:
the go family fdgo (DelayPro), fdanti (DelayAnti), delaygo
(MemoryPro), delayanti (MemoryAnti), reactgo (ReactGo), reactanti
(ReactAnti); the decision-making family delaydm1, delaydm2,
contextdelaydm1, contextdelaydm2, multidelaydm; and the matching family
dmsgo, dmsnogo, dmcgo, dmcnogo. The single-task, two-task and
pretraining experiments train subsets of the go family.
Run every script from the repository root; data is written to the top-level
data directories. Experiment scripts import the shared library in core/
through a small _bootstrap.py shim.
| Path | Contents |
|---|---|
core/ |
Shared library: model (mpn), tasks (mpn_tasks), training (net_helpers, networks), clustering, seed selection, utilities |
one_task/ |
Single-task training, analysis, pipeline |
two_task/ |
Two-task training, analysis, pipeline |
multiple_task/ |
Multi-task training, clustering analysis, lesions, state space, pipeline |
pretrain/ |
Pretraining → post-training transfer experiment and analysis |
flex_task/ |
Flexible-task (RNN/MPN) training and analysis |
paper_plot.py |
Publication figures |
docs/ |
Per-analysis documentation (see Documentation) |
test/ |
Unit tests (python -m unittest discover -s test) |
Python 3.9+, PyTorch (CUDA optional), NumPy, SciPy, scikit-learn, Matplotlib, seaborn, h5py, hdf5plugin, scienceplots. Training and the lesion experiments need a GPU node; the lesion pipeline takes about five hours per run.
Each experiment family follows train → analyze → figures. Every training
script adds N_TRIALS new seeds of the setting configured in the script and
never overwrites a saved run, so rerun it to grow a cohort; --n-trials N
changes the count and --seed S [S ...] trains explicit seeds.
# Multi-task: train, then cluster, lesion and plot the checkpoints that have
# no pipeline outputs yet
python multiple_task/multiple_task.py
python multiple_task/run_pipeline.py --feature L21e4 --skip-completed
# state-space geometry of the saved runs and its link to the lesions
python multiple_task/state_space_shift.py
python multiple_task/state_space_structure.py
python multiple_task/state_space_vs_lesion.py
# share of synapse tuning explained by the endpoint neurons' clusters
python multiple_task/synapse_variance_decomposition.py --seed all --feature L21e4
# Single- and two-task (train + analyze)
python one_task/run_one_task_pipeline.py
python two_task/run_two_task_pipeline.py
# Pretraining transfer: train, then analyze the cohort of one feature label
python pretrain/pretraining.py
python pretrain/pretraining_analysis.py --feature L21e3lam0.9
python pretrain/pretraining_post.py --feature L21e3lam0.9
# Paper figures: everything, one mode, or one figure
python paper_plot.py
python paper_plot.py multiple_tasks lesion --multitask-l2 1e4
python paper_plot.py pretraining --pretraining-lambda 0.9
python paper_plot.py --only state_space_vs_lesionCheckpoints and result files share one identifier, which the analysis scripts parse to find a run's files:
{task}_seed{seed}_{feature}+hidden{hidden}+batch{batch}{accfeature}
# e.g. everything_seed749_L21e4+hidden300+batch128+angle
| Directory | Contents |
|---|---|
multiple_tasks/ |
Multi-task checkpoints, recorded activity, learning curves |
multiple_tasks_analysis/ |
Clustering results and per-run analysis figures |
multiple_tasks_perf/, multiple_tasks_norm/ |
Lesion results and lesion plots |
state_space/, two_in_multiples/ |
State-space and sibling-task geometry caches |
onetask/, twotasks/ |
Single- and two-task runs and analyses |
pretraining/, pretraining_analysis/ |
Pretraining runs and analysis caches |
paper_plot/ |
Publication figures |
- docs/multitask_analysis.md: clustering caches, the lesion figures and their statistics, the modulation clustering variants.
- docs/state_space.md: state-space and sibling-task geometry analyses and figures.
- docs/pretraining.md: the transfer experiment, its figure options and output directories.
- SCHEME.md: timing convention and figure color scheme.
Parts of this codebase were written with the assistance of Claude Code.