Self-distillation-based spiking neural network training method

Abstract

The invention discloses a spiking neural network training method based on self-distillation. The method comprises the following steps: initializing a network structure of a spiking neural network (SNN) and LIF neuron parameters; fusing the outputs of all the time steps according to the fusion weight to obtain a global soft target as a teacher model; combining the output of each time step with historical output according to a fusion weight, and obtaining a local soft target of each time step as a student model; the method comprises the following steps: mapping a global soft target and a local soft target to a shared feature space, calculating self-distillation loss through KL divergence, calculating cross entropy loss for an image sample, integrating the two losses, calculating to obtain a total loss function, updating network parameters of an SNN through back propagation, and applying the trained SNN to an image classification task. According to the method, the coarse-grained time step fusion problem of an existing SNN training method can be improved, and the training precision of the spiking neural network is effectively improved.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of deep learning, and particularly relates to a pulse neural network training method based on self-distillation. BACKGROUND

[0002] In recent years, super-large artificial neural network (ANN) models such as ChatGPT and Gemini have made breakthroughs in language understanding, image generation and multi-modal reasoning tasks. Large-scale network models have become a future development trend. However, behind the success of these models is a huge amount of computing power and energy consumption: for example, the single training of GPT-3 with 175 billion parameters consumes 128 million kilowatt-hours of electricity, emits about 500 tons of carbon dioxide, and still occupies a large amount of GPU resources during the reasoning phase. This phenomenon reveals the bottleneck of the traditional dense matrix computing paradigm in hardware efficiency, and also makes “green AI” an urgent demand of the academic and industrial communities.

[0003] Spiking Neural Network (SNN) is considered a potential solution to the energy consumption problem because it uses discrete pulses of biological neurons “firing-resting” to transmit information. The event-driven nature of SNN means that computation is triggered only when the membrane potential crosses the threshold, naturally having highly sparse spatiotemporal activation; at the same time, the pulse sequence encodes information in the time domain, and can complete perception and decision-making tasks with very low power consumption. Experiments show that on neuromorphic chips (such as Loihi-2, SpiNNaker-2), SNN networks can complete image classification or path planning with only tens of milliwatts, which is in sharp contrast to traditional GPUs, which often consume tens to hundreds of watts for the same task.

[0004] Although SNN has significant energy advantages, its training still faces two major challenges. First, the discontinuity of the pulse activation function causes the error to be unable to be directly back-propagated, leading to gradient vanishing or explosion, limiting the depth of the network and the complexity of the task. To solve this problem, the academic community has proposed three main methods: one is to imitate the STDP (Spike-Timing-Dependent Plasticity) rule of biological plasticity, but it lacks a global error signal and can only work on shallow networks and simple tasks; the second is to use a surrogate gradient to approximate the pulse derivative, which alleviates the non-differentiable problem, but the accuracy of deep networks is still difficult to compete with the corresponding ANN; the third is ANN-to-SNN conversion, which maps high-precision ANN weights to the pulse domain to obtain higher accuracy, but it needs hundreds or even thousands of time steps to work stably, which offsets the natural advantage of SNN in energy consumption. Second, some studies use external ANN teachers to distill knowledge to SNN to improve its performance. However, the coexistence of double models destroys the low-power ecology of pure SNN and adds additional storage and computing burden to the deployment end.

[0005] Self-distillation technology has shown the potential to significantly improve model representation ability without an external teacher by introducing multiple identical or different models, fusing their outputs as a teacher model, and training each model as a student model in a self-distillation training framework in recent years. Unlike continuous activation ANN, the time pulse sequence of SNN naturally contains multi-scale information: the cumulative pulse in a long time window reflects the global decision trend, and the early pulse determines the real-time response speed.

[0006] In SNN models, a natural self-distillation framework is provided, so if the time domain knowledge can be extracted within the same SNN and the corresponding teacher and student models are found, the accuracy gap between SNN and artificial neural network can be narrowed. This is a very worthwhile research problem, so an innovative and effective self-distillation-based pulse neural network training framework is needed to make the trained model have higher accuracy closer to artificial neural network and make the training process more efficient and low-power. SUMMARY

[0007] The present application aims to overcome the shortcomings of the prior art and provide a self-distillation-based pulse neural network training method.

[0008] The purpose of the present application is achieved by the following technical solution: a self-distillation-based pulse neural network training method, comprising the following steps: (1) Obtain an image classification task dataset and perform data preprocessing including image normalization on the images; (2) Select the spiking neural network to be trained, initialize its network structure and LIF neuron parameters, and set learnable cross-temporal fusion weights for the simulation time step T; (3) After the forward propagation of the spiking neural network is completed, the output of time step T is fused according to the cross-temporal fusion weight to obtain the global soft target of the whole time step, which is used as the self-distilled teacher model to provide basic knowledge; (4) For each time step t, its output is fused with the output of its historical time steps according to the cross-temporal fusion weight to obtain the local soft target of the current time step t, which is used as the self-distilled student model. (5) The spiking neural network is trained using the preprocessed image classification task dataset. During the training process, the global soft target output by the teacher model and the local soft target output by the student model are mapped to the shared feature space through a unified fully connected layer and converted into probability through the softmax activation function to obtain the class probability output of the teacher model and the student model. (6) Calculate the self-distillation loss and cross-entropy loss based on the class probability outputs of the teacher model and the student model and the corresponding true labels, and sum them by weight to obtain the total loss function; (7) With minimizing the total loss function as the optimization objective, backpropagation is performed on the total loss function to update the network parameters of the spiking neural network in the student model, and the trained spiking neural network is obtained for image classification tasks.

[0009] Furthermore, the data preprocessing includes the following steps: The content of each channel of the image is normalized; according to the preset time window size T, the image is processed by pulse spike conversion to construct an image sequence of length T.

[0010] Furthermore, step (3) specifically includes: After the forward propagation of the spiking neural network is completed, the pulse information output at T time steps with a time window of [1,T] is weighted and fused according to the cross-temporal fusion weights to obtain the global soft target at all time steps. Using it as a self-distilling teacher model Used to provide basic knowledge.

[0011] Furthermore, step (4) specifically includes: For each time step t, the pulse information output at the current time step t is fused with the pulse information output at its historical time steps [1, t-1], i.e., the pulse information output at the t time steps with a time window of [1, t], according to the cross-temporal fusion weights, to obtain the local soft target at the current time step t. It is used as a student model for self-distillation. .

[0012] Further, the class probability outputs of the teacher model and the student model are denoted as wherein, represents the global soft target is the class probability output of the teacher model for class c, represents the local soft target at time step t is the class probability output of the student model for class c, is the spike count of class c within the time window [1, T], is the spike count of class c within the time window [1, t], and the denominator is the normalization.

[0013] Further, the self-distillation loss is obtained by calculating the KL divergence, and the calculation formula is: wherein, is the self-distillation loss, represents the global soft target is the class probability output of the teacher model for class c, represents the local soft target at time step t is the class probability output of the student model for class c.

[0014] Further, the cross-entropy loss is obtained by calculating the sum of the global cross-entropy loss and the local cross-entropy loss, and the calculation formula is: wherein, is the cross-entropy loss, is the global cross-entropy loss, is the local cross-entropy loss at time step t, represents the onehot true label of class c.

[0015] Further, the calculation formula of the total loss function is: wherein, is the total loss function, and respectively represent the weights of the cross-entropy loss and the self-distillation loss .

[0016] The beneficial effects of the present application are: the self-distillation based pulse neural network training method can make the trained SNN model have higher accuracy closer to artificial neural networks without introducing external neural networks, and the training process is more efficient and low-consumption, further improving the accuracy of the SNN model; the outputs of each time step of the SNN model are mapped to a unified feature space, and the joint loss of the self-distillation loss and the cross-entropy loss can further mine the information between the time steps, so that the low time step obtains more information of the high time step, further improving the accuracy of the SNN model, and the accuracy of the low time step is further improved, reducing the dependence of the pulse neural network on the time step during inference. BRIEF DESCRIPTION OF DRAWINGS

[0017] Figure 1 is a flowchart of the self-distillation based pulse neural network training method of the present application; Figure 2 is a calculation method schematic diagram of the self-distillation loss of the present application; Figure 3 is a framework schematic diagram of the self-distillation based pulse neural network training method of the present application. DETAILED DESCRIPTION

[0018] Hereinafter, exemplary embodiments will be described in detail with reference to the accompanying drawings. In the following description, unless otherwise expressly specified, the same numbers in different drawings represent the same or similar elements. The embodiments described in the following exemplary embodiments are not representative of all embodiments consistent with the present application. Rather, they are merely examples of devices and methods consistent with some aspects of the present application as detailed in the appended claims. It should be understood that the above general description and the following detailed description are merely exemplary and explanatory, and are not limiting to the present application.

[0019] The terms used in the present application are merely for the purpose of describing specific embodiments, and are not intended to limit the present application. The singular forms "a", "an" and "the" used in the present application and the appended claims are also intended to include the plural forms, unless the context clearly indicates otherwise. It should also be understood that the term "and / or" used herein means and includes any or all possible combinations of one or more associated listed items.

[0020] ​It should be understood that, although the terms first, second, third, etc. can be employed in this application to describe various information, these information should not be limited to these terms. These terms are only used to differentiate one piece of information from another piece of information. For example, without departing from the scope of the present application, the first information can also be referred to as the second information, and similarly, the second information can also be referred to as the first information. Depending on the context, the word "if" as used herein can be interpreted to mean "when" or "upon" or "in response to determining". Also, the terms "comprise", "comprising", or any other variant thereof are intended to cover non-exclusive inclusion, so that processes, methods, articles, or equipment including a series of elements include not only those elements, but also other elements not explicitly listed, or other elements inherent to such processes, methods, articles, or equipment. Without more limitations, the element defined by the phrase "comprising a" does not exclude the presence of additional identical elements in the process, method, article, or equipment including the element.

[0021] The present application will be described in detail below with reference to the accompanying drawings. The features in the following examples and embodiments can be combined with each other without conflict.

[0022] Referring to Figure 1 and Figure 3 The self-distillation-based pulse neural network training method of the present application specifically includes the following steps: (1) Obtain an image classification task data set, such as CIFAR-10, CIFAR-100, ImageNet, etc., and perform data preprocessing on the image, including image normalization.

[0023] Further, the data preprocessing includes the following processes: image normalization of the contents of each channel of the image; according to a preset time window size T, the image is processed by spike conversion (Spike Convert) to construct an image sequence with a length of T.

[0024] (2) Select a pulse neural network (SNN) to be trained, initialize the network structure and the parameters of the LIF neuron, and set the cross-time domain fusion weight of the simulation time step T as a learnable parameter. The initialization content includes but is not limited to: the number of network layers of the SNN, the structure of each layer of the network, the feature channel size, the convolution kernel size, the initial value of each layer weight, and the initial parameter value of the LIF neuron.

[0025] It should be understood that the signal of each unit in the SNN is a pulse sequence information with a length of a time step T, which is the simulation time step T, and the simulation time step T is a tunable hyperparameter.

[0026] (3) After the forward propagation of the SNN ends, the outputs of time steps T are fused according to the cross-time domain fusion weight to obtain a global soft target of all time steps, which is used as a teacher model of self-distillation to provide basic knowledge.

[0027] Specifically, after the forward propagation of the SNN ends, the pulse information of T time steps in a time window [1, T] is fused according to the cross-time domain fusion weight to obtain a global soft target of all time steps , which is used as a teacher model of self-distillation to provide basic knowledge.

[0028] (4) For each time step t, the output thereof is fused with the outputs of historical time steps thereof according to the cross-time domain fusion weight to obtain a local soft target of the current time step t, which is used as a student model of self-distillation to provide a synchronous consistency constraint for the current time step t.

[0029] Specifically, for each time step t, the pulse information of the current time step t is fused with the pulse information of historical time steps [1, t-1] thereof, that is, the pulse information of t time steps in a time window [1, t] is fused according to the cross-time domain fusion weight to obtain a local soft target of the current time step t , which is used as a student model of self-distillation .

[0030] It should be noted that, since , finally, T local soft targets can be obtained.

[0031] (5) The SNN is trained by using the image classification task data set after data preprocessing, and in the training process of the SNN, the global soft target output by the teacher model and the local soft target output by the student model are mapped to a shared feature (logits) space through a unified fully connected layer (FC) to align the semantic scale, and are converted into probabilities through a softmax activation function to obtain the class probability outputs of the teacher model and the student model.

[0032] Specifically, the SNN is trained with the data pre-processed image classification task dataset. In the training process of the SNN, the output of each stage time step, i.e. the global soft target output by the teacher model and the local soft target output by the student model, is mapped to a shared feature space through a unified fully connected layer, in which the specific semantics caused by the network structure and time sequence of the SNN are eliminated, so that the feature semantics can be aligned. Then the global soft target and the local soft target are converted into corresponding probabilities through a softmax activation function, and the final class probability output of the teacher model and the final class probability output of the student model are obtained. The class probability output corresponding to the global soft target of the teacher model is used as the basis knowledge for the self-distillation SNN training. From the perspective of the teacher, the final class probability output of the teacher model is directly used as the knowledge source of the local soft target of each time step t.

[0033] Further, the class probability outputs of the teacher model and the student model are respectively represented as: In the formula, represents the global soft target is the class probability output of the teacher model of class c, represents the local soft target of time step t is the class probability output of the student model of class c, is the pulse spike count of class c in the time window [1, T], is the pulse spike count of class c in the time window [1, t], and the denominator is the normalization. It should be noted that since the SNN generates discontinuous pulse activation signals, when calculating the corresponding class probability output, the pulse spike of the pulse activation signal needs to be converted into a probability metric, which is achieved by counting the pulse spike in a specified time window and normalizing it, as shown in the above formula.

[0034] (6) The self-distillation loss and the cross-entropy loss are calculated according to the class probability outputs of the teacher model and the student model and the corresponding real labels (labels), and the weighted sum of the two is obtained to obtain the total loss function.

[0035] Further, as shown in Figure 2 and Figure 3 , the self-distillation loss and the cross-entropy loss are calculated for the pulse neural network with a simulation time step of T. The self-distillation loss is obtained by calculating the KL divergence, and the calculation formula is: In the formula, is the self-distillation loss, represents a global soft target is the class probability output of the teacher model for class c, represents a local soft target at time step t is the class probability output of the student model for class c. The cross-entropy loss is obtained by calculating the sum of the global cross-entropy loss and the local cross-entropy loss, and the calculation formula is: In the formula, is the cross-entropy loss, is the global cross-entropy loss, is the local cross-entropy loss at time step t, represents the onehot real label of class c.

[0036] Further, the calculation formula of the total loss function is: In the formula, is the total loss function, and respectively represent the weights of the cross-entropy loss and the self-distillation loss for balancing the cross-entropy loss and the self-distillation loss.

[0037] (7) Taking minimizing the total loss function as the optimization goal, the back propagation is performed on the total loss function to update the network parameters of the SNN in the student model, including the network weight, the LIF neuron, etc., so that the multi-round training of the SNN is completed, and the trained SNN is obtained for the image classification task. After calculating the total loss function of each batch, the network parameters of the SNN are updated according to the back propagation algorithm, so that the training of the pulse neural network model is completed.

[0038] It should be noted that the training method described in the present application can also be used to train various pulse neural networks, including pulse convolutional neural networks and transformer-like pulse neural networks, etc. The finally trained SNN can also be used for various downstream tasks, such as image classification tasks, image recognition tasks, etc.

[0039] In summary, the training method of the present application does not need to introduce an external artificial neural network model, but uses the output of the same SNN at different time scales to construct "global soft target" and "local soft target", extracts deeper information between time steps through knowledge distillation, maps the pulse features of different time steps to the same feature space, eliminates the inconsistency of the fusion output semantic space of the pulse neural network model at different time steps, aligns the feature semantics, specifically, a new method of applying knowledge distillation to the pulse neural network and a new loss function are introduced to facilitate this process, combined with cross-entropy and self-distillation loss for joint optimization, thereby updating the network parameters, extracting information between time steps, better training the pulse neural network, and thus narrowing the gap in accuracy between the pulse neural network and the artificial neural network. The specific process includes: initializing the pulse neural network and LIF neuron parameters, fusing all simulation step pulses to generate a global soft target, generating a local soft target by fusing historical pulses by time step, calculating and weighting the integrated loss in the shared feature space, and updating the learnable parameters by using alternative gradient-BPTT (backpropagation through time) backpropagation. The present application can significantly shorten the effective simulation step number, reduce the pulse trigger number and energy consumption in the inference stage, while improving the SNN accuracy in image classification, multi-modal and other tasks, and is suitable for convolutional and Transformer-like pulse networks, providing a general and efficient training scheme for low-power and green artificial intelligence. From the global time step, the present application further extracts information between time steps, which can improve the accuracy of the pulse neural network; from the local time step, it can significantly reduce the dependence of the pulse neural network on long simulation time steps, allowing it to have the accuracy of high simulation time steps in a lower time step network model.

[0040] The above examples are only used to illustrate the technical solutions of the present application, and not to limit them; although the present application has been described in detail with reference to the foregoing examples, those skilled in the art should understand that they can modify the technical solutions described in the foregoing examples, or make equivalent substitutions for part of the technical features; and these modifications or substitutions do not make the essence of the corresponding technical solutions deviate from the spirit and scope of the technical solutions of the embodiments of the present application.

Claims

1. A training method for a spiking neural network based on self-distillation, characterized in that, Includes the following steps: (1) Obtain the image classification task dataset and perform data preprocessing on the images, including image normalization; (2) Select the spiking neural network to be trained, initialize its network structure and LIF neuron parameters, and set learnable cross-temporal fusion weights for the simulation time step T; (3) After the forward propagation of the spiking neural network is completed, the output of time step T is fused according to the cross-temporal fusion weight to obtain the global soft target of the whole time step, which is used as the self-distilled teacher model to provide basic knowledge; (4) For each time step t, its output is fused with the output of its historical time steps according to the cross-temporal fusion weight to obtain the local soft target of the current time step t, which is used as the self-distilled student model. (5) The spiking neural network is trained using the preprocessed image classification task dataset. During the training process, the global soft target output by the teacher model and the local soft target output by the student model are mapped to the shared feature space through a unified fully connected layer and converted into probability through the softmax activation function to obtain the class probability output of the teacher model and the student model. (6) Calculate the self-distillation loss and cross-entropy loss based on the class probability outputs of the teacher model and the student model and the corresponding true labels, and sum them by weight to obtain the total loss function; (7) With minimizing the total loss function as the optimization objective, backpropagation is performed on the total loss function to update the network parameters of the spiking neural network in the student model, and the trained spiking neural network is obtained for image classification tasks.

2. The spiking neural network training method based on self-distillation according to claim 1, characterized in that, The data preprocessing includes the following steps: The content of each channel of the image is normalized; according to the preset time window size T, the image is processed by pulse spike conversion to construct an image sequence of length T.

3. The spiking neural network training method based on self-distillation according to claim 1, characterized in that, Step (3) specifically includes: After the forward propagation of the spiking neural network is completed, the pulse information output at T time steps with a time window of [1,T] is weighted and fused according to the cross-temporal fusion weights to obtain the global soft target at all time steps. Using it as a self-distilling teacher model Used to provide basic knowledge.

4. The spiking neural network training method based on self-distillation according to claim 1, characterized in that, Step (4) specifically includes: For each time step t, the pulse information output at the current time step t is fused with the pulse information output at its historical time steps [1, t-1], i.e., the pulse information output at the t time steps with a time window of [1, t], according to the cross-temporal fusion weights, to obtain the local soft target at the current time step t. It is used as a student model for self-distillation. .

5. The spiking neural network training method based on self-distillation according to claim 1, characterized in that, The class probability outputs of the teacher model and the student model are respectively expressed as: In the formula, Represents a global soft target For the class probability output of the teacher model of class c, Represents the local soft target at time step t. The output of the class probability for the student model of class c. Count the pulse spikes of category c within the time window [1,T]. The count is the number of pulse spikes of category c within the time window [1,t], and the denominator is the normalization.

6. The spiking neural network training method based on self-distillation according to claim 1, characterized in that, The self-distillation loss is obtained by calculating the KL divergence, and the calculation formula is as follows: In the formula, This is due to losses from self-distillation. Represents a global soft target For the class probability output of the teacher model of class c, Represents the local soft target at time step t. This is the class probability output for the student model of class c.

7. The spiking neural network training method based on self-distillation according to claim 6, characterized in that, The cross-entropy loss is obtained by calculating the sum of the global cross-entropy loss and the local cross-entropy loss, and its calculation formula is as follows: In the formula, For cross-entropy loss, For global cross-entropy loss, The local cross-entropy loss at time step t. This represents the one-hot true label for category c.

8. The spiking neural network training method based on self-distillation according to claim 7, characterized in that, The formula for calculating the total loss function is as follows: In the formula, For the total loss function, and These represent the cross-entropy loss, respectively. and self-distillation loss The weight.

Images