This repository contains the code for our ICML 2024 paper:

"Sign Gradient Descent-based Neuronal Dynamics: ANN-to-SNN Conversion Beyond ReLU Network"

Spiking neural network (SNN) is studied in multidisciplinary domains to (i) enable order-of-magnitudes energy-efficient AI inference and (ii) computationally simulate neuro-scientific mechanisms. The lack of discrete theory obstructs the practical application of SNN by limiting its performance and nonlinearity support. We present a new optimization-theoretic perspective of the discrete dynamics of spiking neurons. We prove that a discrete dynamical system of simple integrate-and-fire models approximates the sub-gradient method over unconstrained optimization problems. We practically extend our theory to introduce a novel sign gradient descent (signGD)-based neuronal dynamics that can (i) approximate diverse nonlinearities beyond ReLU and (ii) advance ANN-to-SNN conversion performance in low time steps. Experiments on large-scale datasets show that our technique achieves (i) state-of-the-art performance in ANN-to-SNN conversion and (ii) is the first to convert new DNN architectures, e.g., ConvNext, MLP-Mixer, and ResMLP.

• Prerequisites
• Installation
  • Using VSCode Dev Containers
  • Manual Installation
• Usage
• Project Structure
• License
• Contact
• Citation
• Acknowledgements

If you have the Dev Containers extension installed in VSCode:

1. Open the Command Palette:

   • Press Ctrl+Shift+P (or Cmd+Shift+P on macOS).

2. Select:

   • "Dev Containers: Rebuild and Reopen in Container".

3. Wait for Setup:

   • VSCode will rebuild and reopen the project in a containerized environment with all dependencies installed.

If you are not using VSCode:

1. Install Python 3.10:

   • Download and install Python 3.10 from the official website.

2. Install PyTorch 2.0.1 with CUDA:

   • Visit PyTorch's official website and select the appropriate command for your system.
   • Example installation command:

     pip install torch==2.0.1+cu117 torchvision==0.15.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117

3. Clone the Repository:

   git clone https://github.com/snuhcs/snn_signgd
   cd snn_signgd

4. Install Required Packages:

   pip install setuptools
   pip install -e .

We provide a minimal example to demonstrate how to use the code.

1. Navigate to the examples Directory:

   cd examples

2. Run the Minimal Example Script:

   python minimal_example.py

   This script showcases the basic usage of our signGD-based neuronal dynamics for ANN-to-SNN conversion.

The repository is organized as follows:

.
├── data/
├── docs/
├── examples/
├── experiments/
├── src/packages/
│   └── snn_signgd/
├── tests/
├── .devcontainer/
├── LICENSE
├── README.md
├── requirements.txt
├── setup.py

• data/: Stores datasets used in the project, such as CIFAR10.

• docs/: Contains documentation resources, including images and presentation slides that illustrate the concepts and results of the project.

• examples/: Provides example usage scripts for the package. This includes minimal examples to get started and scripts demonstrating how to apply the package to specific datasets like CIFAR10.

• experiments/: Houses experimental configurations, logs, resources, scripts, and source code for running experiments.

  • config/: Configuration files for datasets, models, neuronal dynamics, and optimizers.

    • datasets_and_models/: Configuration files for various datasets (e.g., CIFAR10, CIFAR100, ImageNet) and models.

    • neuronal_dynamics/: Configurations for different neuronal dynamics implementations, such as sign gradient descent and subgradient methods.

    • optimizers/: Configuration files for different optimizers like Adam, SGD, and their variants.

  • logs/: Stores log files generated during experiments.

  • resources/: Contains resources like checkpoints and pretrained models.

  • samples/: Placeholder directory for sample outputs or data generated during experiments.

  • scripts/: Includes scripts for training and evaluating models.

    • train.py: Script to train models using specified configurations.

    • evaluate.py: Script to evaluate trained models on test datasets.

  • src/: Source code specific to experiments, including datasets, models, and tasks.

    • dataset/: Contains dataset loaders and preprocessing scripts.

    • model/: Includes model definitions and architectures used in experiments.

    • task/: Defines tasks for downstream applications like image classification.

• src/: Contains the core code of the package that is installable.

  • packages/: The main package directory that includes all installable modules.

    • snn_signgd/: Core implementation of the sign gradient descent-based neuronal dynamics.

      • conversion.py: Tools for converting ANNs to SNNs using the proposed method.

      • core/: Core modules essential for the package functionality.

      • graph_functional/: Utilities for graph transformations and pattern matching in neural networks.

      • neuronal_dynamical_system/: Implements various neuronal dynamics, including generalizations for different activation functions.

• tests/: Contains unit tests to ensure the correctness and reliability of the package. (To be added)

• .devcontainer/: Configuration files for setting up a development container in VSCode, including the Dockerfile and devcontainer.json.

• LICENSE: The MIT License under which the project is distributed.

• README.md: The main README file providing an overview of the project, installation instructions, usage examples, and other essential information.

• requirements.txt: Lists all Python dependencies required to run the project.

• setup.py: Installation script that allows the package to be installed via pip.

This project is licensed under the MIT License - see the LICENSE file for details.

For questions or inquiries, please contact:

• Hyunseok Oh
• Email: [email protected]
• Affiliation: Seoul National University

If you find this code useful in your research, please cite our paper:

@article{oh2024sign,
  title={Sign Gradient Descent-based Neuronal Dynamics: ANN-to-SNN Conversion Beyond ReLU Network},
  author={Oh, Hyunseok and Lee, Youngki},
  journal={arXiv preprint arXiv:2407.01645},
  year={2024}
}