Summary
The authors propose SpikeGPT, which is a hybrid variant of RWKV architecture that employs spiking linear transformations and some float-point operations. In a rough estimate, SpikeGPT has an energy efficiency advantage of about 32 times over vanilla GPT. In terms of performance, SpikeGPT outperforms the LSTM backbone and is comparable to some simplified variants of the Transformer, but falls behind the vanilla Transformer.
Strengths
1. The paper is easy to follow.
2. The authors replace the linear layer in RWKV, which has the highest computational overhead, with a spiking layer. As a result, SpikeGPT is about 32 times more energy efficient compared to the vanilla GPT.
Weaknesses
1. SpikeGPT introduces some float-point operations including float-point multiplication, division, and exponentiation. This makes SpikeGPT different from traditional spiking neural networks. Although these operations are not dominant in terms of computational overhead, this part of the operation will not be sparse and event-driven, which puts a higher demand on possible application scenarios that must support this hybrid computational paradigm.
2. The normalization used in SpikeGPT is unclear. As shown in the left subplot of Fig. 1, Add&Norm is used in SpikeGPT, but it is not specified in the text whether it is a layernorm or a batchnorm (or other normalization). Batchnorm is generally used for SNNs because it can be merged into linear or convolutional layers, whereas layernorms, which are widely used in NLP, cannot be merged, which results in extra floating-point multiplications and is not applicable to typical SNNs. Based on the code given by the authors, I would guess that the authors use layernorm, but the authors don't provide any analysis of the computational overhead of normalization and the effect of layer normalization on spike inputs in the main text. In addition, Eq. 3 and Eq. 4 do not mention the normalization.
3. Estimates of energy efficiency in the main text are rough. First, the energy consumption estimates in Table 1 assume that all spiking neurons have the same firing rate of 0.15. The authors do not explain how the 0.15 is derived here, but it is safe to assume that this is only a rough estimate, since the firing rate of neurons in different layers and even different channels of the network should be different. Second, the energy consumption estimates in Table 1 do not seem to take into account the differences in energy consumption for different kinds of floating-point operations. In spiking RWKV, the element-wise float-point, division, and exponentiation (including sigmoid) should have different energy consumption with the MAC.
4. Missing further analysis of $\rm {ReLU}^2$ activation. In page 7, line 5, section 3.7, the authors state that "While $\rm {ReLU}^2$ activations are not binary, they induce dynamical sparsity." However, the authors do not give an estimate of this sparsity. In addition, the energy estimates in Table 1 do not take into account the energy consumption of $\rm {ReLU}^2$, nor do they use the sparsity of $\rm {ReLU}^2$ with $\rm E_{MAC}$ to estimate the energy consumption of the linear transformation $\bf M_s$.
5. The contribution on improving RWKV is unclear. The authors only review vanilla self-attention in the main text, but instead of reviewing the RWKV structure they put this part in the appendix. In addition, the authors do not clearly indicate which parts are the authors' main contributions and which are RWKV contributions in the main text. Therefore I need to carefully compare the content about RWKV in the appendix to determine the author's main contribution. As far as I understand from the main text, the author's contribution to the improvement of the RWKV structure is as follows:
- They add a spiking neuron layer before the linear layer (or equally, add a spiking neuron layer after block output) to reduce the computational overhead of the linear layer.
- They modify the mechanism of the positional weight decay in RWKV.
Questions
1. (Weakness 2) What kind of normalization is used in SpikeGPT? Please add it to Eq. 3 and Eq. 4.
2. In page 5, line 11, inline equation $R=X[t]M_R$, $K=X[t]M_K$, and $V=X[t]M_V$. Do you mean $R[t]=X[t]M_R$, $K[t]=X[t]M_K$, and $V[t]=X[t]M_V$?