Skip to content

AaronSpieler/elmneuron

Repository files navigation

The ELM: an Efficient and Expressive Phenomenological Neuron Model Can Solve Long-Horizon Tasks

This repository features a minimal implementation of the (Branch) Expressive Leaky Memory neuron in PyTorch. Notebooks to train and evaluate on NeuronIO are provided, as well as pre-trained models of various sizes.

The ELM neuron

Installation:

  1. Create the conda environment with conda env create -f elm_env.yml
  2. Once installed, activate the environment with conda activate elm_env

Models:

The models folder contains various sized Branch-ELM neuron models pre-trained on NeuronIO.

$d_m$ 1 2 3 5 7 10 15 20 25 30 40 50 75 100
#params 4601 4708 4823 5077 5363 5852 6827 8002 9377 10952 14702 19252 34127 54002
AUC 0.9437 0.9582 0.9558 0.9757 0.9827 0.9878 0.9915 0.9922 0.9926 0.9929 0.9934 0.9934 0.9938 0.9935

We also include a best effort trained ELM neuron that achieves 0.9946 AUC with $d_m=100$.

Notebooks:

  • train_elm_on_shd.ipynb: train an ELM neuron on SHD or SHD-Adding Dataset.
  • train_elm_on_neuronio.ipynb: train an ELM neuron on NeuronIO Dataset.
  • eval_elm_on_neuronio.ipynb: evaluate provided models on the NeuronIO Dataset.
  • neuronio_train_script: script to train an ELM neuron on NeuronIO Dataset.

Code:

The src folder contains the implementation and training/evaluation utilities.

  • expressive_leaky_memory_neuron.py: the implementation of the ELM model.
  • neuronio: files related to visualising, training and evaluating on the NeuronIO dataset.
  • shd: files related to downloading, training and evaluating on the Spiking Heidelberg Digits (SHD) dataset, and its SHD-Adding version.

Note: the PyTorch implementation seems to be about 2x slower than the jax version unfortunately.

Dataset:

Running the NeuronIO related code requires downloading the dataset first (~115GB).

Running the SHD related code is possible without seperately downloading the dataset (~0.5GB).

  • The small SHD daset will automaticall be downloaded upon running the related notebook.
  • A dataloader for the introduced SHD-Adding dataset is provided in /src/shd/shd_data_loader.py
  • For more information onf SHD, please checkout the following website: spiking-heidelberg-datasets-shd

Running the LRA training/evaluation is not provided at the moment.

  • To download the dataset, we recommend to checkout the following repository: mega
  • For the input preprocessing, please refer to our preprint.

Citation:

If you like what you find, and use an ELM variant or the SHD-Adding dataset, please consider citing us:

[1] Spieler, A., Rahaman, N., Martius, G., Schölkopf, B., & Levina, A. (2023). The ELM Neuron: an Efficient and Expressive Cortical Neuron Model Can Solve Long-Horizon Tasks. arXiv preprint arXiv:2306.16922.