A dynamic pruning method of visual transformer
By combining token pruning and feature sparsity to remove redundancy, the high storage and computational costs of visual Transformer models on resource-constrained devices are addressed, achieving efficient model compression and accuracy preservation.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- SOUTH CHINA UNIV OF TECH
- Filing Date
- 2023-06-13
- Publication Date
- 2026-05-12
AI Technical Summary
The deployment of existing visual Transformer models on resource-constrained devices is limited by high storage and floating-point computation costs. Furthermore, the existing dynamic token pruning method is singular, resulting in limited reduction in computational load, and important tokens may be mistakenly discarded.
The model redundancy is removed by a combination of token pruning and feature sparsity. The token matrix depends on the input features of the current layer, and the tokens that are mistakenly discarded can be recovered in the subsequent encoder. The redundant channels are dynamically pruned through the token decision matrix and the feature sparsity module.
It effectively reduces the computational and storage requirements of the model while maintaining model accuracy, enabling efficient deployment on resource-constrained devices.
Smart Images
Figure CN116933859B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of deep learning, and more particularly to a dynamic pruning method for a visual Transformer. Background Technology
[0002] Deep learning is widely used in many real-world applications, such as object detection, image recognition, and semantic segmentation. Since its introduction into computer vision, the Transformer model has demonstrated excellent performance in various visual tasks; however, its high storage and floating-point computation costs limit its deployment on resource-constrained devices. For a visual model, carefully designing a compression scheme for the model architecture to achieve adaptable deployment on resource-constrained target hardware is crucial for the research and deployment of deep neural networks.
[0003] Dynamic pruning plays a crucial role in deploying models on resource-constrained platforms. It can dynamically remove model redundancy based on input, thereby significantly reducing the computational and storage requirements of the network. However, the challenge lies in integrating multiple pruning techniques to prune model redundancy from multiple perspectives while maintaining relatively high model accuracy.
[0004] A dynamic token pruning technique for visual Transformers (DynamicViT Efficient VisionTransformers with Dynamic Token Sparsification) progressively prunes less important tokens in a token sequence to reduce the model's floating-point computation. However, this scheme has a relatively simple pruning method, only pruning the token sequence, resulting in limited reduction in model computation. Furthermore, the token decision matrix in this scheme depends on the current decision matrix for updates; once a token is discarded, it cannot be recovered, which may lead to important tokens being mistakenly discarded in shallow encoder modules. Summary of the Invention
[0005] The purpose of this invention is to provide a dynamic pruning method for visual Transformers to address the problems existing in the prior art. This invention employs a joint approach of token pruning and feature sparsity to simultaneously trim redundant tokens and channels in the model, which can more effectively reduce the floating-point computation of the model and save the computational resources required by the model. Furthermore, the calculation of the token matrix in this invention depends only on the input features of the current layer; tokens mistakenly discarded in shallow encoder modules can be recovered in subsequent encoder modules.
[0006] The present invention is achieved by at least one of the following technical solutions.
[0007] A dynamic pruning method for visual Transformers includes the following steps:
[0008] S1. Construct an image dataset;
[0009] S2. Establish a visual Transformer network model by stacking twelve encoder blocks after the position embedding module. Insert a token pruning module after the third, sixth, and ninth encoder blocks to prune the token sequence, and embed a channel sparsity module in each encoder block to sparsify the feature channels.
[0010] In the location embedding module, the input image is... Reprocessed into a two-dimensional token sequence Where (H, W) is the resolution of the original input image, K is the number of channels of the original input image, N is the length of the token sequence, and M is the embedding dimension of the token. At the same time, a learnable positional encoding vector is constructed to preserve the spatial positional information between the input image blocks. Explicit positional encoding is performed by adding the positional embedding vector to the input embedding vector. Finally, the obtained token sequence is used as the input of the subsequent encoder block.
[0011] S3. Train the visual Transformer network model, optimize the weights and structural parameters of the visual Transformer network model, and prune model redundancy during training. The pruning object of token pruning is the token sequence, and the pruning object of feature sparse pruning is the channel parameters of the fully connected layer.
[0012] S4. Select the sub-token sequence with the highest global information score and the channel with the highest importance weight according to the corresponding ratio, and reconstruct the lightweight network.
[0013] Furthermore, the image data in the image dataset is split into a training set and a test set; the training set is used to train and optimize the visual Transformer network model in step S3, and the test set is used to evaluate the model performance.
[0014] Furthermore, in step S3, token pruning constraints and feature sparsity constraints are added to the classification validation loss function during model search. For a single image, the loss function is:
[0015] L = l c +λ token l token +λ d l d
[0016] Among them l c It is cross-entropy loss, l token It is a token pruning regularization term, λtoken It is the weight of the token pruning regularization term, l d It is the feature sparsity regularization term, λ d It is the weight of the feature sparsity regularization term.
[0017] During model training, the average loss of all images in the image dataset is taken. As training loss:
[0018]
[0019] In the formula, T is the number of images in the image dataset, and L t Let represent the training loss for the t-th image.
[0020] Furthermore, step S3 simultaneously prunes redundant tokens and redundant channels.
[0021] Furthermore, for token pruning, a token decision matrix is first defined.
[0022]
[0023] Where 0 indicates discarding tokens, 1 indicates retaining tokens, and N is the number of tokens; during initialization, the token decision matrix has all values of 1, which means that all tokens are retained initially; during training, the latest decision matrix is calculated in real time in the token pruning module based on the output of the previous encoder block, and the corresponding redundant tokens are discarded.
[0024] Furthermore, each encoder block includes a multi-head self-attention module and a multilayer perceptron module.
[0025] Furthermore, the multilayer perceptron module includes multiple normalization layers, fully connected layers, and nonlinear activation layers.
[0026] Furthermore, the attention matrix of the multi-head self-attention module is calculated as follows:
[0027]
[0028] Where A is the attention matrix, M is the embedding dimension, Q is the query vector, and K is the key vector.
[0029] Furthermore, in the token pruning module, the local feature z of the token is first calculated. local :
[0030]
[0031] Where MLP1(.) is the feature mapping function corresponding to the multilayer perceptron, x is the input of the multilayer perceptron module, and C is specifically taken as C / 2. Let N represent a real matrix with N rows and C′ columns;
[0032] Compute the global feature z of the token glo :
[0033]
[0034] Where MLP2 is the feature mapping function corresponding to the multilayer perceptron with the same structure as MLP1, D is the token decision matrix, and Agg is the token information aggregation function, used to aggregate token context information. Let C' represent a real vector of length C'.
[0035] Furthermore, the token information aggregation function Agg(·,·) employs an average pooling function:
[0036]
[0037] Where D i That is, the decision value for the i-th token in the token decision matrix D, where u is the information aggregation object, u i This is the value at the index corresponding to the information aggregation object.
[0038] Compared with existing technologies, the beneficial effects of the present invention are as follows:
[0039] The present invention provides a dynamic pruning method for visual Transformers, which uses the model compression method described above to dynamically prune the visual Transformer network model. The token pruning and feature sparsity work together to effectively compress the model. It discards redundant tokens with low information content in the image and dynamically skips the calculation of redundant channels, thus combining the advantages of both token pruning and channel sparsity compression methods. Attached Figure Description
[0040] Figure 1 This is a schematic diagram of the dynamic pruning method described in this invention;
[0041] Figure 2 This is a schematic diagram of the model framework of the model compression method described in this invention;
[0042] Figure 3 This is a schematic diagram illustrating the sparse channel features of the model compression method described in this invention. Detailed Implementation
[0043] To enable those skilled in the art to better understand the present invention, the invention will be further described in detail below with reference to the accompanying drawings and specific embodiments. Obviously, the described embodiments are merely some embodiments of the present invention, and not all embodiments. All other embodiments obtained by those skilled in the art based on the embodiments of the present invention without creative effort are within the scope of protection of the present invention.
[0044] Token pruning and channel sparsity methods can effectively compress the model, but using a single pruning strategy is insufficient for removing redundancy. Therefore, this invention employs a combined approach, simultaneously pruning redundant tokens and channels, and combining the advantages of both token pruning and feature sparsity. This results in a lightweight network that meets the resource requirements of the target deployment devices, such as... Figures 1-3 As shown, the dynamic pruning method for the visual Transformer of this invention comprises the following specific steps:
[0045] S1. Divide the original dataset into training and test sets, and perform preprocessing on the images in the dataset, such as padding, cropping, flipping, and normalization. Train and optimize the model's weights and structural parameters on the training set.
[0046] S2. Establish a visual Transformer network model. The visual Transformer model is constructed by stacking twelve encoder blocks after the location embedding module. Each encoder block consists of a multi-head self-attention module and a multilayer perceptron module. At the same time, a token pruning module is inserted after the third, sixth and ninth encoder blocks to prune the token sequence. A channel sparsity module is embedded in each encoder block to sparsify the feature channels. The token pruning target is the token sequence, and the feature sparsity pruning target is the channels of the fully connected layer.
[0047] In the location embedding module, the input image is... Reprocessed into a two-dimensional token sequence Where (H, W) is the resolution of the original input image, K is the number of channels in the original input image, N is the length of the token sequence, and M is the embedding dimension of the tokens. A learnable location encoding vector is constructed to preserve the spatial location information between input image patches. This location embedding vector is added to the input embedding vector through simple summation. Finally, the resulting token sequence is used as the input for subsequent encoder blocks. For the token pruning module, the importance of the token sequence is evaluated, and the latest token decision matrix D is generated based on this evaluation. Then, the token decision matrix is used to prune the token sequence, eliminating the influence of pruned tokens on the classification results.
[0048] Specifically, first, a token decision matrix D is initialized:
[0049] D∈{0,1} N
[0050] Where 0 indicates discarding tokens, 1 indicates retaining tokens, and N is the number of tokens; during initialization, the token decision matrix has all values of 1, meaning that all tokens are initially retained. During training, the token pruning module calculates the latest decision matrix in real time based on the output of the previous encoder block and discards corresponding redundant tokens to reduce the computational load of the model.
[0051] In the token pruning module, the local feature z of the token is first calculated. local :
[0052]
[0053] Where MLP1 is the feature mapping function corresponding to the multilayer perceptron, which specifically consists of a series of normalization layers, a fully connected layer, and a nonlinear activation layer. x is the input of this layer, N is the token sequence length, and C′ is specifically taken as C / 2. This represents a real matrix with N rows and C′ columns.
[0054] Similarly, calculate the global feature z of the token. global :
[0055]
[0056] Where MLP2 is the feature mapping function corresponding to the multilayer perceptron with the same structure as MLP1, C is specifically taken as C / 2, D is the token decision matrix, and Agg is the token information aggregation function, used to aggregate token context information. This represents a real vector of length C.
[0057] As one example, in its implementation, the token information aggregation function Agg adopts an average pool function, namely:
[0058]
[0059] Where D i That is, the decision value for the i-th token in the token decision matrix D, where u is the information aggregation object, u i This is the value at the index corresponding to the information aggregation object.
[0060] Local information encodes the local information of a token, while global information contains the context of the entire image. This allows us to obtain the local-global embedding of the token.
[0061]
[0062] Where zi It is the token local-global embedding matrix, z local It is the token local score matrix, z globa This is the token global score matrix, where N is the token sequence length. The local-global embedding matrix is then fed into another multilayer perceptron (downMLP) to obtain the token discard-retention matrix.
[0063]
[0064] In downMLP, three fully connected layers are used to progressively reduce the dimensionality of the information, and two non-linear activation functions are embedded in the middle for feature extraction. This represents a real matrix with N rows and 2 columns.
[0065] Next, π is sampled to obtain the decision matrix. Since direct sampling of π would result in non-differentiability, hindering end-to-end model training, the Gumbel-Softmax method is used to sample π. Gumbel-Softmax is a normalization function, and its specific calculation method is as follows:
[0066]
[0067] Where y i Let U be the normalized weight matrix, τ be the normalized temperature, U(0,1) represent a uniform distribution between 0 and 1, and o i o j For noise, π i π j The value represents the corresponding token decision probability, and `st` indicates that the code is restricted to the formula below. During training, the normalized temperature `τ` is set to decay exponentially with each epoch, so that the sampled output after training approximates a one-hot code.
[0068] Using the above formula, a differentiable discrete sampling formula is constructed by replacing the conventional argmax with a differentiable Softmax function. During forward training, a differentiable discrete one-hot code is obtained by hard sampling, which can approximate the gradient of the Softmax function during backpropagation, thus enabling automatic gradient updates to optimize the model.
[0069] The decision matrix generated during training is usually unstructured. If the token decision matrix D is directly used... i Discarding tokens at the value of 0 will result in inconsistent token counts within the same batch, hindering parallel computation. Furthermore, simply setting the discarded token to zero does not eliminate its influence on other tokens in the self-attention mechanism. This is because the attention matrix in the self-attention mechanism is calculated as follows:
[0070]
[0071] Where A is the attention matrix, M is the embedding dimension, Q is the query vector, and K is the key vector, both of which are calculated from the input x through a fully connected layer.
[0072] Tokens set to zero still affect other tokens during the Softmax operation. Therefore, a special attention masking strategy is employed, as follows:
[0073]
[0074]
[0075]
[0076] Where P is the attention weight matrix before normalization, P ij P ik That is, the value at the corresponding index. This represents a real matrix with N rows and N columns, where G is the attention masking matrix. ij This is used to indicate whether the j-th token affects the i-th token. Specifically, when i = j, its value is 1, representing that the token will affect itself; when i ≠ j, its value corresponds to the value D in the token decision matrix. j To maintain consistency, G ik G ik That is, the value at the corresponding index. That is, the attention matrix obtained through calculation.
[0077] Based on the above strategy, an attention calculation mechanism with a masking strategy is designed, which can mask the contribution of discarded tokens in the self-attention mechanism while keeping the shape of the output attention matrix N×N.
[0078] For features that are sparse, the input features are first downsampled using the average pooling function:
[0079]
[0080] Where avg(x) l-1 ) is the feature vector obtained after downsampling, x l-1 This is the output of the previous module, where M is the embedding dimension and x is the number of elements. n，m This represents the value of the m-th channel in the n-th token. Let D represent a real vector of length D.
[0081] Then, a learning gate is used to extract features and activate the weights to obtain the feature channel weights of the current layer. The learning gate consists of a fully connected layer and a ReLU activation function.
[0082] g l =ReLU(avg(x) l-1 W1+b1)
[0083] Where g l Let W1 represent the channel weight vector, W1 be the weight matrix of the fully connected layer, and b1 be the bias vector of the fully connected layer.
[0084] Next, a weighted competition selection strategy is adopted for the feature weights:
[0085] π l (x l-1 ) = W αM (g l )
[0086] Here, W represents the selection strategy based on weighted competition. Specifically, the channel weights are sorted, and αM channels with the highest importance based on their retention rate are retained, where α is the channel retention rate and M is the embedding dimension. The remaining low-weight channels are set to 0, and their calculations are skipped in subsequent calculations, thus reducing the computational cost of the model.
[0087] Finally, the classification tokens in the output of the last encoder block are mapped to the dataset category dimension as the model's classification result output:
[0088]
[0089] Where p is the classification result, clsMLP includes one normalization layer and one fully connected layer, and x cls It is the category token in the output of the last-level module, where C is the number of categories. Let C represent a real vector of length C.
[0090] S3. Train the visual Transformer model and update and optimize the network's weights and structural parameters, while pruning redundant tokens and channels.
[0091] For token pruning, as an example during the inference phase, the Softmax function is directly used to predict the token retention probability:
[0092]
[0093] in The token retention probability matrix is defined, and π is the token discard-retention matrix. Then, the corresponding tokens are discarded according to the threshold.
[0094] Determining the Loss Function: The dynamic pruning method of the visual Transformer in this invention can be applied to visual image classification tasks. To improve the model's classification performance, cross-entropy loss is first added to guide the optimization of the model's classification prediction results. Secondly, it is desired that the number of tokens retained in each token pruning stage is as consistent as possible with the expected ratio; therefore, token pruning constraints are added. Finally, L1 regularization constraints are added to the channel tokens to guide their sparsity. The design of the loss function is illustrated below using a single image as an example. Specifically, for the model classification loss, cross-entropy loss is used:
[0095]
[0096] In the formula l c This is the cross-entropy loss value, where C is the number of classes, and y is the cross-entropy loss value. i p represents the true label of an image sample in the image dataset with respect to the i-th category. Its value is 0 or 1, where 1 indicates that the image belongs to that category and 0 indicates that the image does not belong to that category. i This represents the probability that the visual Transformer network model predicts the label of a sample as the i-th category.
[0097] For token pruning, as many redundant tokens as possible are removed; therefore, an L1 regularization constraint is added to the token decision matrix:
[0098]
[0099] In the formula l token This refers to the L1 regularization loss of token pruning, where S is the number of token pruning stages and N represents the length of the token sequence. This represents the pruning decision value of the nth token in the pruning phase of the sth token.
[0100] To minimize the impact of pruned channels on the model, L1 regularization constraints are added to the feature channel weights to enhance their sparsity.
[0101]
[0102] Among them l d That is, the L1 regularization loss for sparse channels, where L represents the sparse layer depth and g l This represents the corresponding channel weight coefficient.
[0103] The complete loss function is a weighted sum of the above losses:
[0104] L = l c +λ token l token +λ d l d .
[0105] Among them l c It is cross-entropy loss, l token It is a token pruning regularization term, λ token It is the weight of the token pruning regularization term, l d It is the feature sparsity regularization term, λ d It is the weight of the feature sparsity regularization term.
[0106] During model training, the average loss of all images in the image dataset is taken. As training loss:
[0107]
[0108] In the formula, T is the number of images in the image dataset, and L t Let represent the training loss for the t-th image.
[0109] S4. Select the sub-token sequence with the highest global information score and the channel with the largest importance weight according to the corresponding ratio, and reconstruct the lightweight network model. The reconstructed lightweight network model can effectively reduce the floating-point computation without significantly affecting the image classification performance of the model.
[0110] The present invention provides a dynamic pruning method for visual Transformers, which optimizes the visual Transformer model using the model compression method described above.
[0111] The preferred embodiments of the present invention disclosed above are merely illustrative of the invention. These preferred embodiments do not exhaustively describe all details, nor do they limit the invention to the specific implementations described. Clearly, many modifications and variations can be made based on the content of this specification. This specification selects and specifically describes these embodiments to better explain the principles and practical applications of the invention, thereby enabling those skilled in the art to better understand and utilize the invention. The invention is limited only by the claims and their full scope and equivalents.
Claims
1. A dynamic pruning method for a visual Transformer, characterized in that, Includes the following steps: S1. Construct an image dataset; S2. Establish a visual Transformer network model by stacking twelve encoder blocks after the position embedding module. Insert a token pruning module after the third, sixth, and ninth encoder blocks to prune the token sequence, and embed a channel sparsity module in each encoder block to sparsify the feature channels. In the location embedding module, the input image is... Reprocessed into a two-dimensional token sequence Where (H,W) is the resolution of the original input image, K is the number of channels of the original input image, N is the length of the token sequence, and M is the embedding dimension of the token. At the same time, a learnable positional encoding vector is constructed to preserve the spatial positional information between the input image blocks. Explicit positional encoding is performed by adding the positional embedding vector to the input embedding vector. Finally, the obtained token sequence is used as the input of the subsequent encoder block. S3. Train the visual Transformer network model, optimize the weights and structural parameters of the visual Transformer network model, and prune model redundancy during training. The pruning object of token pruning is the token sequence, and the pruning object of feature sparse pruning is the channel parameters of the fully connected layer. S4. Select the sub-token sequence with the highest global information score and the channel with the highest importance weight according to the corresponding ratio, and reconstruct the lightweight network.
2. The dynamic pruning method for visual Transformer according to claim 1, characterized in that, The image data in the image dataset is split into a training set and a test set; the training set is used to train and optimize the visual Transformer network model in step S3, and the test set is used to evaluate the model performance.
3. The dynamic pruning method for visual Transformer according to claim 2, characterized in that, In step S3, during model search, token pruning constraints and feature sparsity constraints are added to the classification validation loss function. For a single image, the loss function is: L=l c +λ token l token +λ d l d Among them l c It is cross-entropy loss, l token It is a token pruning regularization term, λ token It is the weight of the token pruning regularization term, l d It is the feature sparsity regularization term, λ d It is the weight of the feature sparsity regularization term; During model training, the average loss of all images in the image dataset is taken. As training loss: In the formula, T is the number of images in the image dataset, and L t Let represent the training loss for the t-th image.
4. The dynamic pruning method for visual Transformer according to claim 1, characterized in that, Step S3 simultaneously prunes redundant tokens and redundant channels.
5. The dynamic pruning method for visual Transformer according to claim 1, characterized in that, For token pruning, first define a token decision matrix. Where 0 indicates discarding the token, 1 indicates keeping the token, and N is the number of tokens; During initialization, the token decision matrix has all values of 1, which means that all tokens are retained initially. During training, the latest decision matrix is calculated in real time based on the output of the previous encoder block in the token pruning module, and the corresponding redundant tokens are discarded.
6. The dynamic pruning method for visual Transformer according to claim 1, characterized in that, Each encoder block includes a multi-head self-attention module and a multilayer perceptron module.
7. The dynamic pruning method for visual Transformer according to claim 6, characterized in that, The multilayer perceptron module includes multiple normalization layers, fully connected layers, and nonlinear activation layers.
8. The dynamic pruning method for visual Transformer according to claim 6, characterized in that, The attention matrix of a multi-head self-attention module is calculated as follows: Where A is the attention matrix, M is the embedding dimension, Q is the query vector, and K is the key vector.
9. The dynamic pruning method for a visual Transformer according to any one of claims 1 to 8, characterized in that, In the token pruning module, the local feature z of the token is first calculated. local : Where MLP1(.) is the feature mapping function corresponding to the multilayer perceptron, x is the input of the multilayer perceptron module, and C′ is specifically taken as C / 2. Let N represent a real matrix with N rows and C′ columns; Compute the global feature z of the token global : Where MLP2 is the feature mapping function corresponding to the multilayer perceptron with the same structure as MLP1, D is the token decision matrix, and Agg is the token information aggregation function, used to aggregate token context information. Let C' represent a real vector of length C'.
10. The dynamic pruning method for visual Transformer according to claim 9, characterized in that, The token aggregation function Agg(·,·) uses an average pooling function: Where D i That is, the decision value for the i-th token in the token decision matrix D, where u is the information aggregation object, u i This is the value at the index corresponding to the information aggregation object.