Block perception mask evolution method for efficient NM sparse training

Abstract

The invention discloses a block sensing mask evolution method for efficient NM sparse training, and relates to the technical field of neural network compression and acceleration. The method comprises the following steps: for an N: M sparse mode, dividing a weight vector into weight blocks with a fixed size of M; a weight index is calculated for each weight block, and then a corresponding binary mask is adaptively obtained according to the weight index; calculating a sparse architecture difference for each group of weight blocks to obtain a block sensing score, and then selecting a high-frequency mask block; the binary mask and the high-frequency mask block are combined to realize joint optimization, a new mask is obtained, the original mask is updated through the new mask, pruning and weight recovery are performed, and the training loss is reduced. According to the block perception mask evolution method for efficient NM sparse training, sparse forward and backward propagation is executed in the whole training process, and the training cost is greatly reduced.

Description

Technical Field

[0001] The present invention relates to the field of neural network compression and acceleration technology, and in particular to a block-aware mask evolution method for efficient N_M sparse training. Background Art

[0002] Traditional network sparsity methods can be roughly divided into two categories: unstructured sparsity and structured sparsity. N:M sparsity, which insists on retaining N out of M continuous weights in deep neural networks (DNNs), has been recently developed and is considered a promising sparsity mode in the research community, supported by NVIDIA Ampere Core. N:M sparsity promotes superior storage and computational efficiency, establishes a perfect harmony between model efficiency and accuracy, and surpasses both unstructured and structured sparsity. For example, 2:4 sparsity can achieve 2x speedup on NVIDIA A100 GPU, while unstructured sparsity may further reduce inference speed at the same level of sparsity. As a pioneering work, ASP [Nvidia. Nvidia a100 tensor coregpu architecture, 2020.] adopted a traditional three-stage workflow, including model pre-training, high-magnitude weight extraction, and network fine-tuning. Subsequently, Zhou et al. proposed a method to train N:M sparse networks from scratch, introducing the sparse refined straight-through estimator (SR-STE). More specifically, N weights with higher magnitude are selected in each forward pass, while all weights are updated in the backward stage, using the STE estimator combined with a uniquely designed sparsity penalty term. Zhang et al. further re-frame N:M sparsity as a combination problem, learning the optimal mask for each N:M block.

[0003] As a pioneering work, ASP chooses to mask those weights with smaller magnitudes, intuitively reducing the output derivatives between dense pre-trained weights and N:M sparse weights. However, the pre-training stage inevitably brings a huge training burden. In the literature, a more popular approach is to perform training-time weight selection by updating all weights. In particular, the straight-through estimator (STE) is used to calculate the gradients of all weights, because the currently removed weights always receive no gradients because their corresponding multiplication masks are 0. By dynamically selecting weights with higher magnitudes, such N:M sparse training can effectively improve model performance without even relying on pre-trained weights.

[0004] Despite recent efforts to further enhance N:M sparse training by adding additional norm constraints on pruned weights or progressive sparsification, an important issue remains that dense backpropagation and weight updates continue to result in large resource consumption, which poses a challenge for resource-limited scenarios. Although these methods are very effective in maintaining the performance of sparse networks, most existing works require dense backpropagation to update all weights to discover the optimal N:M mask, resulting in huge training burden and memory cost. Summary of the invention

[0005] The technical problem to be solved by the present invention is to provide a block-aware mask evolution method for efficient N_M sparse training, which performs sparse forward and backward propagation in the entire training process, greatly reducing the training cost.

[0006] The present invention provides a block-aware mask evolution method for efficient N_M sparse training, comprising:

[0007] Step 1: For the N:M sparse mode, divide the weight vector into weight blocks of fixed size M, where M is the number of weights before sparseness;

[0008] Step 2: Calculate the weight index for each weight block, and then set the mask corresponding to the weight block less than or equal to the index to 0 according to the weight index ranked MN from small to large, and set the rest to 1, to obtain the corresponding binary mask, where N is the number of weights after sparseness;

[0009] Step 3: Calculate the sparse architecture difference for each group of weight blocks to obtain the block perception score, then update the block perception score of each group of weight blocks by exponential moving average, and then select the high-frequency mask block according to the updated block perception score;

[0010] Step 4: Combine the binary mask and the high-frequency mask block to obtain a new mask, update the original mask with the new mask, and perform pruning and weight recovery to reduce training loss.

[0011] The embodiments of the present invention have at least the following technical effects:

[0012] We implement loss-aware mask adaptation and oscillation-aware block selection. Loss-aware mask adaptation effectively reduces training loss by pruning and restoring weights, and oscillation-aware block selection restricts mask modifications to blocks that exhibit high-frequency mask oscillations, thereby stabilizing the N:M training process and ensuring consistent sparsity during the forward and backward propagation stages of the N:M sparse training process, thereby solving the problem of low training efficiency. BRIEF DESCRIPTION OF THE DRAWINGS

[0013] The present invention will be further described below in conjunction with embodiments with reference to the accompanying drawings.

[0014] Figure 1 A schematic diagram of a process of an embodiment of the present invention;

[0015] Figure 2 This is a schematic diagram of a performance comparison result of the method of the embodiment of the present invention and different methods for training N:M sparse on various network architectures;

[0016] Figure 3 This is the second schematic diagram of the performance comparison results of the method of the embodiment of the present invention and different methods for training N:M sparse on various network architectures. DETAILED DESCRIPTION

[0017] The embodiments of the present application provide a block-aware mask evolution method for efficient N_M sparse training, which performs sparse forward and backward propagation in the entire training process, thereby greatly reducing the training cost.

[0018] The technical solution in the embodiment of the present application has the following general idea:

[0019] We provide a method (BAME) to ensure consistent sparsity during the forward and backward propagation phases of the N:M sparse training process, addressing the problem of inefficient training. Unique contributions of BAME include loss-aware mask adaptation, which effectively reduces training loss by pruning and restoring weights, and oscillation-aware block selection, which limits mask modifications to blocks that exhibit high-frequency mask oscillations, thereby stabilizing the N:M training process.

[0020] This embodiment provides a block-aware mask evolution method for efficient N_M sparse training, such as Figure 1 As shown, including:

[0021] Step 1: For the N:M sparse mode, divide the weight vector into weight blocks of fixed size M, where M is the number of weights before sparseness; this can be achieved by the following steps:

[0022] For the weights W = {W1, W2, W3, ..., WL} of a convolutional network with L layers, after determining the sparse number of weights M, the weight vector Wk (k = 1, 2, 3, ..., L) is divided into weight blocks of size M {(Wk, 1), (Wk, 2), (Wk, 3), ...}.

[0023] Step 2: Calculate the weight index for each weight block, and then set the mask corresponding to the weight block less than or equal to the index to 0 according to the weight index ranked MN from small to large, and set the rest to 1, to obtain the corresponding binary mask, where N is the number of weights after sparseness; loss-aware mask adaptation is implemented to effectively reduce training loss by pruning and restoring weights; specifically, the following steps may be included:

[0024] (1) Calculate the weight index: For each weight block in the N:M mode weight vector, calculate the weight index according to the following formula:

[0025]

[0026] The above is the derivation process of the calculation formula. In practical applications, the following formula can be used directly for calculation:

[0027]

[0028] Among them, L represents the loss function of model training, W represents the weight of the convolutional network, k represents the kth weight block, i represents the i-th weight in the kth weight block, η represents the current learning rate, and B represents the mask.

[0029] (2) Generate a binary mask for each weight vector: After obtaining the weight index of each weight block in the weight vector, generate a binary mask according to the following formula:

[0030]

[0031] in, Represents the weight index ranked MN among all weight indexes in the kth weight block. If it is less than or equal to this weight index, it is zero. If it is greater than the MNth weight index, it is 1. Then N weights will be retained. For example, if M is 8 and N is 2, the element ranked 6th in weight index is selected. Then the elements less than or equal to the 6th are all zero, and the elements greater than the 6th are 1.

[0032] Step 3: Calculate the sparse architecture difference for each group of weight blocks to obtain a block perception score, then update the block perception score of each group of weight blocks by exponential moving average, and then select the high-frequency mask block (group number) according to the updated block perception score, so as to realize oscillation-aware block selection, and restrict mask modification to blocks showing high-frequency mask oscillation; specifically, the following steps may be included:

[0033] (1) For each group of weight blocks, calculate the sparse architecture difference (Sparse Architecture Divergence) and obtain the block perception score (SAD score):

[0034]

[0035] Among them, B represents the mask, k represents the kth weight block, c-1 represents the previous training, and c represents the current training.

[0036] (2) Update the block perception score of each group of weighted blocks through exponential moving average:

[0037]

[0038] Among them, O represents the block perception score of each weight block, and γ represents the momentum coefficient of the exponential moving average.

[0039] (3) Select high-frequency mask blocks based on block perception scores:

[0040]

[0041] where β represents the ratio of N:M blocks that limits LMA updates to have the highest oscillation frequency, Represents the perception score of the β×k-ranked block under the current training.

[0042] Step 4: Combine the binary mask and the high-frequency mask block to achieve joint optimization (OBF) to obtain a new mask, update the original mask through the new mask, and perform pruning and recovery weights to reduce training losses.

[0043] The binary mask and the high-frequency mask block are specifically combined using the following formula:

[0044]

[0045] in, represents the current training binary mask, represents the high-frequency mask block currently being trained, Represents the high-frequency mask block of the previous training, Represents the kth weight of the current training. α represents the proportion of important elements in the control selection weight.

[0046] The proposed BAME method is different from the existing N:M method because it performs sparse forward and backward propagation throughout the training process, which greatly reduces the training cost. The performance comparison of the proposed method (BAME) and different methods for training N:M sparse on various network architectures is shown in Figure 2 and Figure 3 , where Top-1Accuarcy represents the Top-1 classification accuracy, Epochs represents the current training, and FLOPs represents the number of floating-point operations, which can be understood as the amount of calculation. It can be used to measure the complexity of the algorithm / model.

[0047] from Figure 2 It can be seen that the BAME method proposed in the present invention is in the leading position in all N:M modes and networks, and has a smaller training burden. In the ResNet-50 model, the Top-1 classification accuracy of the BAME technology in the 2:4 mode is improved by 0.2% compared with LBC (LBC is 77.2%, BAME is 77.4%), proving its advantage over ASP and SR-STE in all sparse modes. Specifically, ASP uses 200 rounds including pre-training and fine-tuning stages to train the Top-1 accuracy of 2:4 sparse ResNet-50 to 76.8%, while the BAME method obtains a significantly higher accuracy of 77.4% using 120 rounds.

[0048] from Figure 3 It can be seen that in the ResNet-32 model, the Top-1 classification accuracy technology of the BAME method in the 2:4 mode is 0.18% higher than that of LBC (LBC is 94.81%, BAME is 94.99%). In the 1:4 and 1:16 modes, the BAME method is also in a leading position. In the MobileNet-V2 model, the Top-1 classification accuracy technology of the BAME method in the 2:4 mode is 0.84% ​​higher than that of LBC (LBC is 92.48%, BAME is 93.32%).

[0049] The present invention realizes loss-aware mask adaptation and oscillation-aware block selection. Loss-aware mask adaptation effectively reduces training loss by pruning and restoring weights. Oscillation-aware block selection restricts mask modification to blocks exhibiting high-frequency mask oscillations, thereby stabilizing the N:M training process and ensuring consistent sparsity in the forward and backward propagation stages of the N:M sparse training process, thereby solving the problem of low training efficiency.

[0050] Although the specific implementation modes of the present invention are described above, those skilled in the art should understand that the specific implementation modes described are only illustrative and are not intended to limit the scope of the present invention. Equivalent modifications and changes made by those skilled in the art in accordance with the spirit of the present invention should be included in the scope of protection of the claims of the present invention.

Claims

1. A block-aware mask evolution method for efficient N_M sparse training, characterized in that include: Step 1: For the N:M sparse mode, divide the weight vector into weight blocks of fixed size M, where M is the number of weights before sparseness; Step 2: Calculate the weight index for each weight block, and then set the mask corresponding to the weight block less than or equal to the index to 0 according to the weight index ranked MN from small to large, and set the rest to 1, to obtain the corresponding binary mask, where N is the number of weights after sparseness; Step 3: Calculate the sparse architecture difference for each group of weight blocks to obtain the block perception score, then update the block perception score of each group of weight blocks by exponential moving average, and then select the high-frequency mask block according to the updated block perception score; Step 4: Combine the binary mask and the high-frequency mask block to obtain a new mask, update the original mask with the new mask, and perform pruning and weight recovery to reduce training loss.

2. The method according to claim 1, characterized in that: Calculate the weight index of each weight block using the following formula: Among them, L represents the loss function of model training, W represents the weight of the convolutional network, k represents the kth weight block, i represents the i-th weight in the kth weight block, η represents the learning rate of the current training, and B represents the mask.

3. The method according to claim 2, characterized in that: The binary mask is generated according to the following formula: in, Represents the weight of the MN ranked from small to large among the ownership in the k-th weight block.

4. The method according to claim 1, characterized in that: In step 3, the calculation formula of the block perception score is: Among them, B represents the mask, k represents the kth weight block, c-1 represents the previous training, and c represents the current training.

5. The method according to claim 4, characterized in that: The block perception score of each group of weight blocks is updated by exponential moving average using the following formula: Where O represents the block-aware score of each weight block, and γ represents the momentum coefficient of the exponential moving average.

6. The method according to claim 5, characterized in that: The following formula is used to select the high-frequency mask block according to the updated block perception score: where β represents the ratio of N:M blocks that limits LMA updates to have the highest oscillation frequency, Represents the perception score of the β×k-ranked block under the current training.

7. The method according to claim 6, characterized in that: The combination of the binary mask and the high frequency mask block specifically adopts the following formula: in, represents the current training binary mask, represents the high-frequency mask block currently being trained, Represents the high-frequency mask block of the previous training, Represents the kth weight of the current training. α represents the proportion of important elements in the control selection weight.

Images