This is the official code for the paper "Memory-enriched computation and learning in spiking neural networks through Hebbian plasticity" ([Abstract] arxiv.org, [PDF] arxiv.org).
You need PyTorch to run this code. We tested it on PyTorch version 1.7.0. Additional dependencies are listed in environment.yml. If you use Conda, run
conda env create --file=environment.ymlto install the required packages and their dependencies.
Below you will find a short description of each task presented in the paper and instructions on how to reproduce the results.
In this task we tested the ability of our model to one-shot memorize associations and to use these associations later when needed. The task requires to form associations between random continuous-valued vectors and integer labels that were sequentially presented to the network.
To start training on the association task, run
python memorizing_associations_task.py [optional arguments]Set the command line argument --sequence_length to set the number of vector-label pairs (in the paper we used
sequences ranging from 1 to 55). Here, the command line argument --num_classes should be set to the same value.
To test the out-of-distribution capability of our model in this task, set --sequence_length < --num_classes and then
test the model with 2 < --sequence_length <= --num_classes.
To evaluate a trained model on the test data set, run:
python memorizing_associations_task.py --resume='PATH_TO_CHECKPOINT_FILE' --evaluate [optional arguments]The optional arguments must be set to the same values as during training, otherwise an error is thrown.
To plot the network activity and the model's output after training, run:
python plot_memorizing_associations_task.py --checkpoint_path='PATH_TO_CHECKPOINT_FILE' [optional arguments]The optional arguments must be set to the same values as during training, otherwise an error is thrown.
Here we applied our model to the problem of 1-shot 5-way classification on the Omniglot data set. We used a CNN as input encoder for the Omniglot images, which was pre-trained using the prototypical loss and then converted into a spiking CNN by using a threshold-balancing algorithm. The checkpoint of this CNN is available here in this repository.
To start training on the Omniglot 1-shot task, run
python omniglot_one_shot_task.py [optional arguments]To evaluate a trained model on the test data set, run:
python omniglot_one_shot_task.py --resume='PATH_TO_CHECKPOINT_FILE' --evaluate [optional arguments]The optional arguments must be set to the same values as during training, otherwise an error is thrown.
To plot the network activity and the model's output after training, run:
python plot_omniglot_one_shot_task.py --checkpoint_path='PATH_TO_CHECKPOINT_FILE' [optional arguments]The optional arguments must be set to the same values as during training, otherwise an error is thrown.
In this task we asked whether Hebbian plasticity can enable SNNs to perform cross-modal associations. We trained our model in an autoencoder-like fashion. We used the FSDD and the MNIST data set in this task. We used two CNNs as input encoder which were pre-trained on FSDD/MNIST classification respectively and then converted into spiking CNNs by using the threshold-balancing algorithm. The checkpoints of these CNNs are available here and here in this repository.
To start training on the cross-modal associations task, run
python cross_modal_associations_task.py [optional arguments]To evaluate a trained model on the test data set, run:
python cross_modal_associations_task.py --resume='PATH_TO_CHECKPOINT_FILE' --evaluate [optional arguments]The optional arguments must be set to the same values as during training, otherwise an error is thrown.
To plot the network activity and the model's output after training, run:
python plot_cross_modal_associations_task.py --checkpoint_path='PATH_TO_CHECKPOINT_FILE' [optional arguments]The optional arguments must be set to the same values as during training, otherwise an error is thrown.
In this task we applied our model to the bAbI data set. We used 10k training examples and trained models on each of the 20 tasks individually.
To start training on bAbI task 1 in the 10k training examples setting, run
python question_answering_task.py [optional arguments]Set the command line argument --task to train on other tasks. To set the synaptic delay in the feedback loop, use the
command line argument --readout_delay (in the paper we have used 1ms and 30ms).
To evaluate a trained model on the test data set, run:
python question_answering_task.py --resume='PATH_TO_CHECKPOINT_FILE' --evaluate [optional arguments]The optional arguments must be set to the same values as during training, otherwise an error is thrown.
To plot the network activity and the model's output after training, run:
python plot_question_answering_task.py --checkpoint_path='PATH_TO_CHECKPOINT_FILE' [optional arguments]The optional arguments must be set to the same values as during training, otherwise an error is thrown.
Here we evaluated our model on an episodic reinforcement learning task. The task is based on the popular children’s game Concentration. We consider a one-player solitaire version of the game. In this version of the game the objective is to find all matching pairs with as few card flips as possible.
To start training on the Concentration game with a deck of four cards, run
python reinforcement_learning_task.py --decay_lr_linearly [optional arguments]Use the command line argument --num_cells to set the number of cards (in the paper we have used 4 and 6). You will
also need to adjust --num_steps (we have used --num_steps=10 for the 4-cards game and --num_steps=100 for the
6-cards game). To use a new deck of cards for each game use --resample_cards.
To evaluate a trained model, run:
python reinforcement_learning_task.py --decay_lr_linearly --resume='PATH_TO_CHECKPOINT_FILE' --evaluate [optional arguments]The optional arguments must be set to the same values as during training, otherwise an error is thrown.
Model training can be distributed across multiple GPUs and multiple nodes (see below for examples). You should always use the NCCL backend for multiprocessing distributed training since it currently provides the best distributed training performance.
python SOME_SCRIPT.py --dist_url 'tcp://127.0.0.1:FREE_PORT' --dist_backend 'nccl' --multiprocessing_distributed --world_size 1 --rank 0 [optional arguments]Node 0:
python SOME_SCRIPT.py --dist_url 'tcp://IP_OF_NODE0:FREE_PORT' --dist_backend 'nccl' --multiprocessing_distributed --world_size 2 --rank 0 [optional arguments]Node 1:
python SOME_SCRIPT.py --dist_url 'tcp://IP_OF_NODE0:FREE_PORT' --dist_backend 'nccl' --multiprocessing_distributed --world_size 2 --rank 1 [optional arguments]If you use this code or models in your research and find it helpful, please cite the following paper:
@article{limbacher2022memory,
title={Memory-enriched computation and learning in spiking neural networks through Hebbian plasticity},
author={Limbacher, Thomas and {\"O}zdenizci, Ozan and Legenstein, Robert},
journal={arXiv preprint arXiv:2205.11276},
year={2022}
}
Authors of this work are affiliated with Graz University of Technology, Institute of Theoretical Computer Science, and Silicon Austria Labs, TU Graz - SAL Dependable Embedded Systems Lab, Graz, Austria. This work was supported by the CHIST-ERA grant CHIST-ERA-18-ACAI-004, by the Austrian Science Fund (FWF) project number I 4670-N (project SMALL), and by the "University SAL Labs" initiative of Silicon Austria Labs (SAL). We thank Wolfgang Maass and Arjun Rao for initial discussions.