Paper: Towards Zero Memory Footprint Spiking Neural Network Training (ICLR-2024)

Reviewer comment (Writing):
The description of the neural node in Section 3.2 is convoluted and hard to follow. The underlying logic of the proposed neural node remains unclear. The motivation behind introducing such a neural node is ambiguous. The rationale for segregating the states into two groups at each time step is not motivated. The necessity of introducing $\hat{V^t}$ (Eqs. 3 and 6) is not clearly justified. Moreover, is $[Y_1,Y_2]$ identical to $[X_1,X_2]$ of the next layer? Is $[\hat{V_1^t},\hat{V_2^t}]$ identical to $[V_1^{t+1},V_2^{t+1}]$? To enhance clarity, I recommend that the authors initially expound upon the LIF model, introduce the components of the model that prompt dissatisfaction, and subsequently describe their modifications based on the LIF model.

Author response:
**Response to *Weakness 1.1 Clarification on spiking neuron Description***: 

The rationale behind introducing the reversible spiking neuron lies in enabling the recalculation of activation values during backpropagation. This approach eliminates the need to store activation values during the forward propagation phase, thereby significantly reducing memory usage. To facilitate the recalculation, our model incorporates a symmetric design, wherein states are divided into two groups.

For a more comprehensive understanding of this process, please refer to **Common Question 1** in the **Global Response** section, where we detail how our reversible spiking neuron substantially reduces memory usage and explain the reversibility of our spiking neuron.

**Response to *Weakness 1.2 Symbol classification***: 

To address your concern about clarifying the variables in our equations, we have provided a detailed explanation of the symbol system employed in our paper within **Common Question 2** of the **Global Response** section. We used the original LIF model as an example to explicitly delineate the meanings of various symbols in our formulas. This illustration helps to distinguish the symbols used in our work from those in the original LIF models.

Additionally, we would like to address the two specific questions you raised in the reviews:
- The symbols $[Y_1,Y_2]$ represent the output of each spiking neuron. After passing through intermediate layers, such as CNN or LayerNorm layers, they transform into the input $[X_1,X_2]$ for the next spiking neuron. 
- The notation $[\hat{V_1^{t}},\hat{V_2^{t}}]$ is the same as $[{V_1^{t}},{V_2^{t}}]$. The hat symbol over $V$ was used to denote that the voltage $V$ undergoes in-place modification during the computation process. We have replaced $\hat{V}$ with $V$ for clarity in our revised version.

Reviewer comment (Theory):
Based on my understanding, $[Y_1, Y_2]$ represents the transmitted signal to other SNN nodes. The issue at hand is that these signals are not binary spikes according to eqs. 2 and 5. However, one of the appealing aspects of SNNs is their utilization of binary events for information processing. In essence, a spiking neuron is active only when it encounters spikes, enabling an event-driven regime and an energy-efficient system. However, the proposed neuron nodes deviate from this advantageous feature by transmitting real-valued signals: neurons are active all the time, and the communication cost is high.

Author response:
**Response to *Weakness 2. signal transmission real-valued vs. binary spikes***: 

Despite the intrinsic appeal of SNNs due to their use of binary events for information processing, the exploration of real-valued outputs in SNNs is a dynamic area of research. A noteworthy example can be found in the work of Guo et al. [1] (ECCV 2022). Their research demonstrates the feasibility of utilizing real-valued spikes in SNNs while maintaining the advantages of binary spikes.

Reviewer comment (Theory):
Why the neural node can achieve layer reversibility during training? Considering the feedforward connection from the l-th layer to the (l+1)-th layer, $x^{l+1}[t] = W * y^l[t]$, where W is not invertible. I am curious about the approach taken to recalculate $y^l[t]$ from $x^{l+1}[t]$. After all, VGG is used as the network architecture, where the weight matrices are not invertible.

Author response:
**Response to *Weakness 3. Explanation for layer reversibility during training***:

Firstly, it is important to highlight that in our model, the transmission of input features $X$ through the spiking neuron does not involve any weight computation. 

Secondly, for layers where weights play a role, such as CNN layers, constructing reversible layers is feasible. A notable illustration of this concept is found in the work of Gomez et al. [2] (NIPS 2017). Since weights are updated only after the completion of the gradient calculation, it allows for the precise recalculation of the original activation values using the unchanged, original weights.

To ensure the integrity of our model, we have rigorously verified the consistency between all inputs in the forward operation and the outputs from the inverse operation. This verification was conducted using `torch.allclose(rtol=1e-06, atol=1e-10)`, confirming the accuracy and reliability of our layer reversibility approach during training.

Reviewer comment (Experiments):
While the issue of high training costs is particularly pronounced for large-scale datasets such as ImageNet, this work exclusively tests the proposed method on small-scale datasets (including Tiny-ImageNet). Consequently, the results may not be as convincing, primarily because the training cost for small-scale tasks is relatively affordable. It would be beneficial to extend the evaluation to more computationally demanding datasets to underscore the true efficacy and efficiency of the proposed method.

Author response:
**Response to *Weakness 4. Training Cost and Scalability***:

Thank you for highlighting this important aspect. In response, we have included an additional experiment in the **Common Question 3** of our **Global Response** section. This experiment specifically targets heavy GPU load tasks, with a focus on training on ImageNet-1k using a single GPU with 24GB memory. The results demonstrate that our method successfully avoids out-of-memory issues even at high batch sizes. We believe this additional data will effectively address your concerns regarding the scalability and computational efficiency of our approach in more demanding scenarios.

Reviewer comment (Experiments):
The approach is only tested on static vision benchmarks that arguably do not require a lot of temporal processing.

Author response:
We agree that evaluating our approach on a mix of both static and dynamic benchmarks is crucial for a comprehensive assessment.  To this end, in the paper we discussed the comparison of our work with the SOTA methods on two static vision datasets, CIFAR10 and CIFAR100, as well as on two dynamic vision datasets, **DVS-CIFAR10** and **DVS128-Gesture**.  These comparative results are comprehensively detailed in **Table 1** of our paper, demonstrating the effectiveness of our approach on varied types of datasets.

Reviewer comment (Experiments):
The paper would become stronger if the RevSNN approach was tested on benchmarks that require long sequences, and where other approaches fail because of out-of-memory issues. An example of such a task could be image classification from sequences of pixels, as described in [1], but other tasks are possible. The goal is to find a setting where only RevSNN manages to learn the task thanks to its low memory requirement.

Author response:
Thank you for highlighting this important aspect. In response, we have included an additional experiment in the **Common Question 3** of our **Global Response** section. This experiment specifically targets heavy GPU load tasks, with a focus on training on ImageNet-1k using a single GPU with 24GB memory. The results demonstrate that our method successfully avoids out-of-memory issues even at high batch sizes. We believe this additional data will effectively address your concerns regarding the scalability and computational efficiency of our approach in more demanding scenarios.

Reviewer comment (Theory):
Does the reversible nature of the RevSNN impose any constraint on the dimensionality of the layers? Specifically, do the dimensions have to be the same for the module to be invertible? If yes, how is this handled in architectures like VGG?

Author response:
RevSNN imposes no constraints on the dimensionality of the layers.  However,  the dimension of a specific layer does have to be the same for the layer to be invertible. 

In a SNN adapted from conventional VGG architectures, a spiking neuron is typically  added after the CNN layers at the construction stage. In pseudocode,  this step is to modify `layer += [nn.Conv2d(in_channels, out_channels]`  to `layer += [nn.Conv2d(in_channels, out_channels], neuron.LIFNode()]`. The reversibility of the two type of layers are discussed below:
* Spiking neuron: Reversibility is inherently ensured due to its identical input and output dimensions
* CNN layers:
For CNN layers with identical input and output dimensions, we can apply the techniques from your suggested reference, *The Reversible Residual Network: Backpropagation Without Storing Activations*, to facilitate reversibility.
For CNN layers with different  input and output dimensions, reversibility cannot be achieved due to the inherent discrepancy in dimensions. However, CNN layers in most VGG models have identical input and output dimensions, leaving this issue with minor affect to the overall memory cost.

Reviewer comment (Presentation):
It is unclear why Figure 4 mentions GNN (Graph Neural Network?), whereas the architectures in the tables are all ConvNets.

Author response:
Thanks a lot for pointing out our typo. We have carefully reviewed and corrected this typo in the revised version of our paper. We hope this amendment addresses your concern and enhances the clarity and accuracy of our work.

Reviewer comment (Presentation):
It would be preferable if the plots in Figure 6 shared the same y-axis, to make comparison easier when varying the timesteps.

Author response:
Thank you for your valuable suggestion. To facilitate easier comparison, we have updated **Figure 6** so that all plots now share the same y-axis. We hope this alignment can enhance clarity and improve the interpretability of the comparisons.

Reviewer comment (Novelty):
The authors might want to consider citing [2] and [3], as they seem related to the approach.

Author response:
We appreciate your valuable suggestion. In line with your recommendation, we have included citations to the two works you suggested in our revised paper.

Reviewer comment (Writing):
The methodology part of this work is not well-presented. There are no intuitions or motivations to demonstrate why we have to design this symmetric forward/backward implementation. The notation system is very chaotic, the authors did not mention what is the input/output/membrane potential, and how you denote traditional LIF nodes' implementation. The figures are just copies of the equation.

Author response:
**Response to *Weakness 1.1: Clarification on Spiking Neural Description***: 

The rationale behind introducing the reversible spiking neuron lies in enabling the recalculation of activation values during backpropagation. This approach eliminates the need to store activation values during the forward propagation phase, thereby significantly reducing memory usage. To facilitate the recalculation, our model incorporates a symmetric design, wherein states are divided into two groups. 

For a more comprehensive understanding of this process, please refer to **Common Question 1** in the **Global Response** section, where we detail how our reversible spiking neuron substantially reduces memory usage and explain the reversibility of our spiking neuron.

**Response to *Weakness 1.2: Symbol classification***: 

To address your concern about clarifying the variables in our equations, we have provided a detailed explanation of the symbol system employed in our paper within **Common Question 2** of the **Global Response** section. We used the original LIF model as an example to explicitly delineate the meanings of various symbols in our formulas. This illustration helps to distinguish the symbols used in our work from those in the original LIF models.

To directly respond to your specific query regarding the definitions of input, output, and membrane potential:
1. **Input potential** at time $t$ is denoted by $V^{t-1}$, representing  the neuron's potential at the previous time step.
2. **Output potential** at time $t$  is denoted by $V^{t}$, indicating the neuron's potential at the current time step.
3. **Membrane potential** at time $t$ is denoted as $M^{t}$, bridging the neuron's states between the previous and the current time step. 

***For Q1:***

SNN models, like the LIF model, require retaining activation values for **each timestep across all layers** of the model, leading to substantial memory usage. Our work aims to eliminate the need to store **any activation values ($X^t$ and $Y^t$) at each time step during training.** To aid in understanding, we reiterate the explanation of the LIF model:

>1. *Charge step*: The calculation of membrane potential at each time step involves storing the previous voltage and input features, leading to a high memory footprint.

>$M^t = V^{t-1} + \frac{1}{\tau}\cdot(X^t - V^{t-1})$

>Here, $V^{t-1}$ represents the neuron’s voltage at time step $t - 1$, and $X^t$ denotes the input features at time step $t$. The term $\frac{1}{\tau}$ acts as a scaling factor. These variables are used to calculate the membrane potential $M^t$ of the spiking neuron at time step $t$.

>2. *Fire step*: The determination of the neuron's output based on the membrane potential and threshold voltage further adds to the memory usage.

>$Y^t = \mathcal{H}(M^t - V_{th})$

>This equation calculates the output at time $t$, denoted as $Y^t$, using the Heaviside function $\mathcal{H}$. It is applied to the difference between the membrane potential $M^t$ and a predefined constant threshold voltage $V_{th}$.

>3. *Reset step*: Updating the voltage for the next time step requires memory to store the current output and membrane potential.

>$V^{t} = (1 - Y^t) \odot M^{t} + Y^t \cdot V_{res}$

>The updated voltage $V^{t}$ is determined by the output $Y^t$ and a constant reset voltage $V_{res}$.

>In this context of the entire process, $V_{t -1}$ is the input voltage, $V_{t}$ is the output voltage and $M^{t}$ is the membrane potential.

Our research is driven by the goal of minimizing this memory footprint. The proposed reversible neuron functions as a layer in the original SNN architecture, with input $X^t$ and output $Y^t$. However, all the $X^t$，$Y^t$ values can be recalculated during the backpropagation process, obviating the need for their storage. While previous works, such as Skipper [1], have utilized checkpoint techniques to reduce memory consumption by storing checkpoints instead of activation values, our approach takes a further step.  It reduces the memory footprint by eliminating the need to save checkpoints, achieving a near-zero memory footprint design for SNNs.

Reviewer comment (Experiments):
This work did not reach the same level of accuracy as the recent SOTA SNN works. For example: TEBN [1], MPBN [2].

Author response:
**Response to *Weakness 2***:

We appreciate your comparison with recent SOTA SNN works. Contrary to the concern raised, our results not only achieve higher accuracy than TEBN but also match the accuracy levels of MPBN.

In our research, we have incorporated our reversible spiking neuron into several existing SOTA SNN methodologies, as detailed in **Table 2**. This was done while maintaining other aspects of their setups unchanged, allowing for a direct comparison of the impact of our innovation on accuracy. Our RevDSR model achieves an accuracy of $95.35$% on CIFAR-10 and $78.21$% on CIFAR-100. Upon reviewing TEBN and MPBN, we found that TEBN attains a maximum accuracy of $94.71$% on CIFAR-10 and $76.41$% on CIFAR-100. Meanwhile, MPBN shows a variable accuracy range from $92.22$% to $96.47$% on CIFAR-10 and $70.79$% to $79.51$% on CIFAR-100 under different settings.

***For Q2:***

Our focus is on enhancing the memory efficiency in SNN models. Typically, SNN models with a higher number of time steps tend to consume more memory. With this prerequisite, we picked **DSR** with T=20 (same as the original paper) and **replaced the original spiking neuron with our reversible spiking neuron** to obtain our **RevDSR**. This experiment was designed to show performance retention and memory enhancement with the reversible spiking neuron.

Due to time constraints during rebuttal, we haven’t integrate our reversible spiking neuron with TEBN. We recognize the importance of a comprehensive evaluation across various SOTA models. In our final version, we plan to include TEBN, following a similar process as with DSR, i.e., replacing its spiking neuron with our reversible spiking neuron. This will provide a broader understanding of the effectiveness of our approach in various SNN architectures.

To address your specific queries:

1. The performance of ReDSR on CIFAR10 with Timstep=2 is  $94.20$%, while the accuracy of TEBN is $94.57$%. 

2. The following table provides detailed information about the neural architectures, time steps, and accuracies in the ImageNet experiments. Our RevSNN model achieves an accuracy of $63.78$%, while TEBN reports an accuracy of  $64.29$%, both with a time step of 4.

|   Method |Architecture | Timestep |   Accuracy(%) | Batch Size 128 | Batch Size 256 | Batch Size 512 |
|:------------:|:-----------:|:-----------:|:-----------:|:----------------:|:----------------:|:----------------:|
| Hybrid Training [2]|ResNet34|250| $61.48%$ | Out-of-Memory | Out-of-Memory | Out-of-Memory |
| Tandem Learning [3]|AlexNet|10|  $50.22%$ | $✓$ | Out-of-Memory | Out-of-Memory |
| OTTT [4]|NF-ResNet-34|6|  $64.16%$ | $✓$ | $✓$ | Out-of-Memory |
| **RevSNN(Ours)** |ResNet-18 |4 | $63.78%$ | $✓$ | $✓$ | $✓$ |

Reviewer comment (Experiments):
Compared to the original SNNs, this new framework increases the training latency. It is okay to have a little bit higher latency, since most memory efficiency methods trade space for time. However, in that case, I hope the authors can demonstrate the real challenging case, i.e. heavy GPU load tasks, like ImageNet-1k training with a single GPU. How much batch size can this method and the traditional method reach and what are the accuracies? These results are missing.

Author response:
**Response to *Weakness 3***:

 Thank you for highlighting this important aspect. In response, we have included an additional experiment in the **Common Question 3** of our **Global Response** section. This experiment specifically targets heavy GPU load tasks, with a focus on training on ImageNet-1k using a single GPU with 24GB memory. The results demonstrate that our method successfully avoids out-of-memory issues even at high batch sizes. We believe this additional data will effectively address your concerns regarding the scalability and computational efficiency of our approach in more demanding scenarios.

***For Q2:***

Our focus is on enhancing the memory efficiency in SNN models. Typically, SNN models with a higher number of time steps tend to consume more memory. With this prerequisite, we picked **DSR** with T=20 (same as the original paper) and **replaced the original spiking neuron with our reversible spiking neuron** to obtain our **RevDSR**. This experiment was designed to show performance retention and memory enhancement with the reversible spiking neuron.

Due to time constraints during rebuttal, we haven’t integrate our reversible spiking neuron with TEBN. We recognize the importance of a comprehensive evaluation across various SOTA models. In our final version, we plan to include TEBN, following a similar process as with DSR, i.e., replacing its spiking neuron with our reversible spiking neuron. This will provide a broader understanding of the effectiveness of our approach in various SNN architectures.

To address your specific queries:

1. The performance of ReDSR on CIFAR10 with Timstep=2 is  $94.20$%, while the accuracy of TEBN is $94.57$%. 

2. The following table provides detailed information about the neural architectures, time steps, and accuracies in the ImageNet experiments. Our RevSNN model achieves an accuracy of $63.78$%, while TEBN reports an accuracy of  $64.29$%, both with a time step of 4.

|   Method |Architecture | Timestep |   Accuracy(%) | Batch Size 128 | Batch Size 256 | Batch Size 512 |
|:------------:|:-----------:|:-----------:|:-----------:|:----------------:|:----------------:|:----------------:|
| Hybrid Training [2]|ResNet34|250| $61.48%$ | Out-of-Memory | Out-of-Memory | Out-of-Memory |
| Tandem Learning [3]|AlexNet|10|  $50.22%$ | $✓$ | Out-of-Memory | Out-of-Memory |
| OTTT [4]|NF-ResNet-34|6|  $64.16%$ | $✓$ | $✓$ | Out-of-Memory |
| **RevSNN(Ours)** |ResNet-18 |4 | $63.78%$ | $✓$ | $✓$ | $✓$ |