A ViT model pruning method based on statistical learning, a pruning device and an image recognition method

Abstract

This invention relates to a statistical learning-based ViT model pruning method, pruning device, and image recognition method. The pruning method includes: using an initial token sequence as the current input token sequence; obtaining the current input token feature sequence through the current encoding layer of a Transformer encoder; determining the number of tokens to be retained in the current encoding layer according to a preset pruning ratio; obtaining the currently retained token feature sequence from the current input token feature sequence based on a statistical learning method, according to the current input token feature sequence and the number of tokens to be retained; and traversing all encoding layers of the Transformer model to obtain the final token features. The pruning method of this invention can reduce the possibility of erroneously deleting important tokens, thereby improving model performance and the accuracy of results for tasks such as image classification.

Description

Technical Field

[0001] This invention relates to the field of computer vision technology, and in particular to a statistical learning-based ViT model pruning method, pruning device, and image recognition method. Background Technology

[0002] Vision Transformer (ViT), as an emerging visual model architecture, has achieved remarkable results in the field of computer vision in recent years. This model has been widely applied to tasks such as image classification, object detection, and image segmentation. Specifically, ViT segments an image into multiple smaller image blocks, performs embedding mapping on each block, and then inputs it into the Transformer architecture for feature extraction and classification. Its powerful parallel computing capabilities and ability to model long-range dependencies give it unique advantages in processing complex image data.

[0003] However, as the model size increases, ViT faces problems such as high computational complexity, large memory consumption, and slow inference speed, which limits its widespread deployment in practical applications, especially in resource-constrained mobile devices or scenarios with high real-time requirements. To address these issues, existing technologies have proposed token pruning techniques, which reduce sequence length by removing unimportant tokens, thereby reducing computational complexity and memory consumption and improving the model's inference speed.

[0004] However, existing token pruning methods have limitations in evaluating token importance and cannot fully reflect the actual contribution of tokens in the model. This may lead to the accidental deletion of some important tokens during the pruning process, which will affect the performance of the model and ultimately result in less than ideal results and lower accuracy in tasks such as image classification, object detection, and image segmentation. Summary of the Invention

[0005] Therefore, the purpose of this invention is to provide a statistical learning-based pruning method for ViT models, which can reduce the possibility of erroneously deleting important tokens, thereby improving model performance and increasing the accuracy of results for tasks such as image classification, object detection, and image segmentation.

[0006] A statistical learning-based pruning method for ViT models includes the following steps:

[0007] Based on the original image, obtain the initial token sequence; The initial token sequence is used as the current input token sequence, and the current input token feature sequence is obtained through the current encoding layer of the Transformer encoder; Based on the preset pruning ratio, determine the number of tokens that need to be retained in the current encoding layer; Based on the current input token feature sequence and the number of tokens that the current encoding layer needs to retain, the current retained token feature sequence is obtained from the current input token feature sequence using a statistical learning method. The next encoding layer is used as the current encoding layer, and the current retained token feature sequence corresponding to the previous encoding layer is used as the current input token sequence. Steps S2-S4 are repeated to traverse all encoding layers of the Transformer model, thereby obtaining the final token features. The step of obtaining the currently retained token feature sequence from the current input token feature sequence includes: Based on the current input Token feature sequence and the number of Tokens that the current encoding layer needs to retain, generate multiple current candidate retained Token combination features; Obtain the total loss function value corresponding to each current candidate retained token combination feature; Select the current candidate retained token combination feature with the minimum total loss function value to obtain the current retained token feature sequence.

[0008] The pruning method described in this invention is based on statistical learning. It obtains the total loss function value corresponding to each candidate token combination feature and selects the minimum total loss function value. Then, based on the candidate token combination feature with the minimum total loss function value, it obtains the current token feature sequence. This method can reduce the possibility of erroneously deleting important tokens, thereby improving model performance and increasing the accuracy of results for tasks such as image classification, object detection, and image segmentation.

[0009] Furthermore, the total loss function value is the sum of the feature reconstruction error and the regularization term.

[0010] Furthermore, the feature reconstruction error is obtained in the following manner: Based on the attention mechanism, the current candidate retained token combination features are aligned to the dimension of the current input token feature sequence to obtain the aligned current candidate retained token combination features; then, the feature reconstruction error between the current input token feature sequence and the aligned current candidate retained token combination features is calculated.

[0011] Further, the aligned current candidate retained token combination features are obtained through the following steps: calculate the attention weight matrix of the current candidate retained token combination features to the current input token feature sequence, and then generate the aligned current candidate retained token combination features by weighted summation.

[0012] Furthermore, the feature reconstruction error between the current input token feature sequence and the aligned current candidate retained token combination features is calculated using L2 norm, L1 norm, or cosine similarity.

[0013] Furthermore, a regularization term is calculated based on the number of tokens that need to be retained in the current encoding layer.

[0014] Further, obtain the feature sequence of the current input token, including: Based on the current input Token sequence, obtain the single-head attention output features of each attention head; Based on the output features of each single-head attention, the multi-head attention encoding features are obtained; Based on the multi-head attention encoding features, the current input token feature sequence is obtained through a feedforward neural network.

[0015] Furthermore, the initial token sequence includes a classification token and multiple image patch embedding vectors, and the current candidate retained token combination features include the classification token and multiple image patch token features.

[0016] The present invention also provides a statistical learning-based ViT model pruning device for implementing the aforementioned pruning method.

[0017] The present invention also provides an image recognition method, which includes the aforementioned pruning method, and further includes: image recognition based on the final token feature.

[0018] To better understand and implement this invention, the following detailed description is provided in conjunction with the accompanying drawings. Attached Figure Description

[0019] Figure 1 This is a schematic diagram of the pruning method of the present invention; Figure 2 This is a schematic diagram of the current input Token feature sequence acquisition process in step S2 of the pruning method of the present invention; Figure 3 This is a schematic diagram of the single-head attention output feature acquisition process in step S2 of the pruning method of the present invention; Figure 4 This is a schematic diagram of the process for obtaining the current retained Token feature sequence in step S4 of the pruning method of the present invention; Figure 5 This is a schematic diagram of the process for obtaining the total loss function value in step S4 of the pruning method of the present invention; Figure 6 This is another schematic diagram of the pruning method of the present invention; Figure 7 This is a schematic diagram of the pruning device of the present invention. Detailed Implementation

[0020] To make the objectives, technical solutions, and advantages of this invention more readily understood by those skilled in the art, the invention will be further described in detail below with reference to the accompanying drawings and specific embodiments.

[0021] refer to Figures 1-5 The statistical learning-based ViT model pruning method of the present invention specifically includes the following steps: S1. Obtain the initial token sequence based on the original image.

[0022] In this step, the image token sequence acquisition module 1 sequentially divides the original image into blocks, embeds and maps them, and combines them into vectors to obtain the initial token sequence.

[0023] In this invention, the original image can be an RGB color image or a grayscale image, and its size can be set according to the actual application scenario, such as 224×224 pixels or 384×384 pixels.

[0024] Block processing involves uniformly dividing the original image into blocks of a preset size, resulting in multiple image blocks. Each image block is a small image region of the same size, containing local information from the original image. The number of image blocks is determined by both the original image size and the block size, calculated using the following formula:

[0025] Where N represents the number of image blocks after the original image is segmented, H and W represent the height and width of the original image, respectively, and P and Q represent the height and width of the image blocks, respectively.

[0026] Each image block obtained after block processing is input into a convolutional layer. Linear projection is performed through the convolution kernel, mapping each image block to an embedding vector of the same dimension. The size of the convolution kernel is the same as the size of the image block, its stride is equal to the size of the image block, and the number of output channels is the dimension of the resulting embedding vector, which can be expressed as:

[0027] in, Let i represent the i-th image patch, and Conv represent the convolution operation. This represents the embedding vector corresponding to the i-th image patch. In a specific embodiment, the dimension of the embedding vector can be set to 768 or 1024.

[0028] After obtaining the embedding vectors corresponding to all image patches, the embedding vectors are arranged according to the original order of the image patches, and then combined to obtain the initial token sequence, as described below:

[0029] in, to This represents the embedding vector corresponding to each image patch. to The vectors have the same dimension. This represents an initial token sequence containing only image patch embedding vectors, with dimensions N×D, where N represents the number of image patches after the original image is segmented, and D represents the dimension of the embedding vector corresponding to the image patch.

[0030] Furthermore, in a specific embodiment where image classification and recognition are subsequently required, after obtaining the embedding vectors of all the aforementioned image patches and arranging the embedding vectors according to the original order of the image patches, a learnable classification token can be added at the starting position to aggregate global information and ultimately use it for the classification task, thereby combining to obtain an initial token sequence. This initial token sequence can then be represented as:

[0031] in, Represents a learnable category token. , to The vectors have the same dimension. This represents the initial token sequence containing categorized tokens, with a dimension of (N+1)×D.

[0032] In this invention, the processing of step S1 realizes the complete transformation process from the original image to the token sequence, providing a data foundation for subsequent Transformer encoder processing.

[0033] S2. Using the initial token sequence as the current input token sequence, obtain the current input token feature sequence through the current encoding layer of the Transformer encoder.

[0034] This step can be performed by the current input Token feature sequence acquisition module 2. In this invention, the initial Token sequence is used as the input to the current encoding layer of the Transformer encoder. At this time, the current encoding layer is the first encoding layer of the Transformer encoder. Specifically, the initial Token sequence is feature extracted through a multi-head attention mechanism and nonlinear transformation is performed using a feedforward neural network to obtain the current input Token feature sequence corresponding to this encoding layer. The specific steps are as follows: S21. Based on the current input Token sequence, obtain the single-head attention output features of each attention head.

[0035] This step is executed through the single-head attention output feature acquisition submodule 21. For each attention head in the multi-attention mechanism, the single-head attention output features of that attention head are obtained using the same method, including: S211. Linearly map the current input Token sequence into a query matrix, a key matrix, and a value matrix, respectively.

[0036] In this step, the current input token sequence undergoes three different linear transformations to obtain the query matrix Q, the key matrix K, and the value matrix V. The three linear transformations are implemented using learnable weight matrices, as shown below:

[0037]

[0038]

[0039] in, , , All are learnable weight matrices, each with a dimension of D× D represents the dimension of the aforementioned embedding vector; For a single attention head, and in multi-head attention mechanisms, it is usually... Set to D divided by the number of attention heads h, i.e. =D / h; This represents the current input token sequence. When the current encoding layer is the first encoding layer, this... This is the initial token sequence. When the current encoding layer is not the first encoding layer, this... This is the currently retained token feature sequence corresponding to the previous encoding layer of the current encoding layer.

[0040] S212. Obtain the attention weight matrix based on the query matrix and the key matrix.

[0041] In this step, the dot product similarity between the query matrix Q and the key matrix K is calculated. Then, this dot product similarity is scaled using a scaling factor and normalized using the Softmax function to obtain the attention weight matrix A. This attention weight matrix reflects the relevance strength between each token and all other tokens. The calculation formula is as follows:

[0042] in, This indicates that matrix multiplication is performed on the transpose of the query matrix and the key matrix to obtain an M×M similarity matrix, where M is the number of tokens in the current input token sequence; The scaling factor is used to prevent the gradient from vanishing due to excessively large dot product results. The Softmax function normalizes the scaled similarity matrix so that the sum of the weights in each row of the matrix is ​​1. The element in the i-th row and j-th column of the attention weight matrix A... This represents the attention weight of the i-th token on the j-th token.

[0043] S213. The attention weight matrix and the value matrix are weighted and summed to obtain the single-head attention output features.

[0044] In this step, the obtained attention weight matrix A and value matrix V are multiplied to obtain the single-head attention output features. Its dimension is M× This step implements a weighted aggregation of the value matrix based on attention weights, enabling each token to incorporate information from other tokens associated with it.

[0045] S22. Based on the output features of each single-head attention, obtain the multi-head attention encoding features.

[0046] This step is executed by the multi-head attention encoding feature acquisition submodule 22, which concatenates all the single-head attention output features to obtain the concatenated feature matrix. Its dimension is M×h That is, M×D; then, a linear transformation layer is used to perform a linear transformation and fusion on the concatenated features to obtain the multi-head attention encoding features. The calculation formula is:

[0047]

[0048] Wherein, Concat represents the concatenation operation, which concatenates the single-head attention output features of h attention heads along the feature dimension; Let be a learnable weight matrix with dimensions D×D; The dimension is M×D.

[0049] S23. Obtain the current input Token feature sequence based on the multi-head attention encoding features.

[0050] This step is performed by the attention encoding feature nonlinear transformation submodule 23, which generates the multi-head attention encoding features. The input is fed into a feedforward neural network for nonlinear transformation. This feedforward neural network consists of two fully connected layers and an activation function, with multi-head attention encoding features. First, the feature dimension is expanded using a first fully connected layer. Then, a non-linear processing is performed using an activation function. Finally, a second fully connected layer restores the feature dimension to the dimension of the multi-head attention-encoded features. The calculation formula is as follows:

[0051]

[0052] in, and These are the weight matrix and bias term of the first fully connected layer, respectively. and These represent the weight matrix and bias term of the second fully connected layer, respectively, and GELU is the activation function of the Gaussian error linear unit. This represents the output of the feedforward neural network, which is the current input token feature sequence obtained by the current coding layer of the Transformer encoder. Its dimension is M×D.

[0053] In step S2, the Transformer encoder's current encoding layer performs feature extraction and nonlinear transformation on the current input Token sequence to obtain the current input Token feature sequence, thus providing a feature basis for subsequent Token pruning decisions.

[0054] S3. Determine the number of tokens that need to be retained in the current encoding layer based on the preset pruning ratio.

[0055] This step is performed by module 3, which obtains the current number of retained tokens, and the pruning ratio... A value between 0 and 1, representing the proportion of tokens that need to be pruned in the current encoding layer.

[0056] With the addition of category tokens, since category tokens are always retained during pruning and do not participate in pruning, the number of tokens to be retained in this step includes category tokens and some image patch tokens. The calculation formula is as follows:

[0057] Where M is the number of tokens in the current input token sequence of the current encoding layer. For pruning ratio, This indicates rounding up, and the max function ensures that at least one token is retained.

[0058] Furthermore, in this invention, when the Transformer encoder includes multiple coding layers, the pruning ratio can be a set, uniform value, or it can be a progressively decreasing ratio sequence designed by an algorithm.

[0059] In this invention, there is no restriction on the order of steps S3 and S2.

[0060] S4. Based on the current input Token feature sequence and the number of Tokens that the current encoding layer needs to retain, obtain the current retained Token feature sequence from the current input Token feature sequence using a statistical learning method.

[0061] In this invention, this step is performed by the currently retained Token feature sequence acquisition module 4, and specifically includes the following steps: S41. Based on the current input Token feature sequence and the number of Tokens that the current encoding layer needs to retain, generate multiple current candidate retained Token combination features.

[0062] This step S41 is executed by the submodule 41 that obtains the current candidate retained token combination features. The number of tokens contained in each candidate retained token combination feature is equal to the number of tokens that the current coding layer needs to retain. Specifically, it includes the following steps: First, select from all tokens in the current input token feature sequence. Each token is retained. Specifically, all image patch token combinations can be obtained through enumeration, heuristic search, or greedy algorithms, thereby generating multiple current candidate token combination features. , The dimension is ×D, where D represents the number of tokens that the current encoding layer needs to retain, and D is the dimension of the embedding vector.

[0063] With the addition of a categorical token, the specific steps are as follows: Identify the categorical tokens in the current input token feature sequence and determine which categorical tokens should be retained. Typically, the categorical token is located at the first position in the sequence, with an index of 0, and it is never pruned in any pruning combination, thus always being retained. Then, after excluding the categorical tokens from the current input token feature sequence, select the appropriate categorical token from the remaining image patch tokens. One token is retained, and all image patch token combinations can be obtained through enumeration, heuristic search, or greedy algorithms. The number of tokens in each image patch token combination is equal to the number of tokens that the current encoding layer needs to retain minus the number of classification tokens. Finally, each image patch token combination is merged with the classification tokens to generate the features of the multiple current candidate retained token combinations. .

[0064] In another embodiment, multiple candidate retained token combination position index combinations can be obtained by first obtaining the position index combination corresponding to each image block token combination, or by obtaining the position index of the classification token and the position index combination corresponding to each image block token combination, and merging the position index of the classification token and the position index combination corresponding to each image block token combination to obtain multiple candidate retained token combination position index combinations; then, according to the candidate retained token combination position index combination, the corresponding token is extracted from the current input token feature sequence to form the current candidate retained token combination feature.

[0065] S42. Obtain the total loss function value corresponding to each current candidate retained token combination feature.

[0066] In this invention, this step is executed by the submodule 42 that obtains the total loss function value, and the specific steps include: S421. Based on the attention mechanism, align the current candidate retained token combination features to the dimension of the current input token feature sequence to obtain the aligned current candidate retained token combination features.

[0067] Due to the characteristics of the current candidate reserved token combination Since the dimensionality of the input token feature sequence is inconsistent, the feature difference cannot be directly calculated. Therefore, an attention mechanism is needed to... Alignment is performed along the dimensions of the current input token feature sequence, as follows: First, calculate the attention weight matrix of the current candidate retained token combination features to the current input token feature sequence. Then, the aligned candidate retained token combination features are generated by weighted summation. The calculation formula is:

[0068]

[0069] in, The current input token feature sequence has dimensions M×D; To preserve the token combination features for the current candidates, the dimension is: ×D, where D is the dimension of the embedding vector; This is the attention weight matrix, with dimensions M× , which represents the attention weight of each token input to the current encoding layer for each retained token; The aligned token feature matrix has dimensions M×D, the same as the dimension of the current input token feature sequence.

[0070] S422. Calculate the feature reconstruction error between the current input Token feature sequence and the aligned current candidate retained Token combination features.

[0071] In this step, the feature reconstruction error is used to measure the feature difference before and after token pruning. It can be calculated using methods such as L2 norm, L1 norm, or cosine similarity. The smaller the value of the feature reconstruction error, the more original information is retained in the pruned token sequence.

[0072] In one specific embodiment, the feature reconstruction error is calculated using the L2 norm, and the calculation formula is as follows:

[0073] Where M is the number of tokens input to the current encoding layer. For the current input Token feature sequence The vector in the i-th row of the vector. Preserve token combination features for the aligned current candidate The vector in the i-th row of the vector. The L2 norm of a vector is the square of its elements, which is the sum of the squares of all its elements. This represents the feature reconstruction error.

[0074] S423. Calculate the regularization term based on the number of tokens that need to be retained in the current encoding layer.

[0075] In this step, the regularization term is used to constrain the number of tokens retained after pruning. The more tokens retained, the larger the regularization term becomes. The calculation formula is as follows:

[0076] in, The number of tokens that the current encoding layer needs to retain. λ is the regularization coefficient used to balance the weight between feature reconstruction error and the number of tokens that need to be retained, and λ is an adjustable hyperparameter, usually set to a small positive number, such as 0.01 or 0.001. This represents the regularization term.

[0077] In this invention, step 423 can also be set before step 422 or even before step S421, so that the regularization term is calculated first and then the feature reconstruction error is obtained.

[0078] S424. Based on the feature reconstruction error and the regularization term, obtain the total loss function value corresponding to each current candidate retained token combination feature.

[0079] In this step, the feature reconstruction error term is added to the regularization term to obtain the total loss function value corresponding to each candidate retained token combination feature. The calculation formula is:

[0080] The total loss function value comprehensively reflects the performance of the candidate token retention combination feature in terms of both information retention degree and token quantity control. The smaller the value, the better the performance of the candidate token retention combination feature in the above two aspects.

[0081] Step S42 of the present invention calculates the attention weight of the retained token to the input token, thereby aligning the pruned feature matrix to the dimension of the current input token feature sequence, which makes it possible to calculate the feature reconstruction error, thereby effectively compensating for the information loss during the pruning process and ensuring that the pruned token feature sequence can retain the important information in the input features to the greatest extent.

[0082] S43. Select the candidate retained token combination features with the smallest total loss function value to obtain the current retained token feature sequence.

[0083] This step is executed by the current retained token feature sequence selection submodule 43, which obtains the minimum total loss function value and its corresponding current candidate retained token combination features. From the current candidate retained token combination features with the minimum total loss function value, the current retained token feature sequence is obtained, which is the optimal token retention scheme determined after statistical learning optimization for the current encoder layer. Its size is the number of tokens that the current encoding layer needs to retain. .

[0084] In another embodiment, the candidate retained token combination position index combination corresponding to the minimum total loss function value can be obtained first. Then, based on this position index combination, the corresponding token is extracted from the current input token feature sequence of the current encoding layer, thereby forming the retained token feature sequence. Specifically, this method involves iterating through each index i in the candidate retained token combination position index combination corresponding to the minimum total loss function value, and extracting the vector of the i-th row from the current input token feature sequence. The elements are arranged according to the order of the indices in the location index combination to form a new feature matrix. This matrix is ​​the currently retained token feature sequence, and its dimension is [missing information]. ×D, where The number of tokens reserved in the current encoding layer, where D is the dimension of the embedding vector.

[0085] In this invention, the currently retained Token feature sequence retains the most important Token information in the current encoding layer, while significantly reducing the sequence length and lowering the computational complexity of subsequent encoding layers.

[0086] Through the above steps S3-S4, a single token importance assessment and pruning decision based on statistical learning methods is realized. It selects the current retained token feature sequence corresponding to the minimum total loss function value from the current input token feature sequence. It contains the feature information of all tokens retained after pruning. Its dimension is smaller than that of the original input sequence, which not only reduces the computational workload of subsequent encoding layers, but also provides a decision basis for token retention and transmission in subsequent encoding layers.

[0087] S5. Take the next encoding layer as the current encoding layer, and take the current retained token feature sequence corresponding to the previous encoding layer as the current input token sequence. Repeat steps S2-S4 to traverse all encoding layers of the Transformer model to obtain the final token features.

[0088] In this invention, this step is performed by the final Token feature acquisition module 5, as referred to... Figure 6 Specifically, the process includes the following steps: determining whether the current layer index is less than the total number of encoding layers L in the Transformer model. This determination is used to determine whether all encoding layers have been processed. If the current layer index is less than the total number of encoding layers, it means that there are still unprocessed encoding layers. The currently retained token feature sequence corresponding to the encoding layer at the current layer index needs to be used as the input of the next encoding layer. Then, proceed to step S2 to continue the loop processing. If the current layer index is equal to the total number of encoding layers, it means that all encoding layers have been processed and the loop ends.

[0089] Specifically, if the current layer index is less than the total number of coding layers, the current layer index is incremented by 1 and initialized. The coding layer corresponding to the initialized current layer index is then used as the current coding layer. Subsequently, in the processing of the next coding layer as the current coding layer, the same processing method as in steps S2-S4 is used. The currently retained token feature sequence output from the previous coding layer is used as the input of the current coding layer. Feature extraction and nonlinear processing are performed on the retained token feature sequence of the previous coding layer to obtain the current input token feature sequence. Then, based on the current input token sequence and the number of tokens to be retained corresponding to the pruning ratio of the current coding layer, the currently retained token feature sequence is obtained from the current input token feature sequence using statistical learning methods. If the current layer index equals the total number of coding layers, the currently retained token feature sequence corresponding to the last coding layer is used as the final token feature. The dimension is ×D, where The final token feature contains the number of tokens retained in the last layer, and D is the dimension of the embedding vector. This final token feature includes the feature information of the classification token and the most important image patch tokens retained after feature extraction and pruning of all encoding layers. The classification token is located at the first position in the final token feature sequence and is used to aggregate global information and perform classification.

[0090] In this invention, the total number of encoding layers L in the Transformer model is a preset hyperparameter of the model, usually set to 12 layers, 24 layers or more, and can be determined according to the model size and task requirements.

[0091] Furthermore, the present invention also provides an image recognition method, which includes the aforementioned steps S1-S5, and further includes step S6, performing image recognition based on the final token feature, and outputting the image recognition result.

[0092] Step S6 can be specifically executed by the image recognition module. Specifically, it involves inputting the final token features into the classification layer, obtaining the predicted probability distribution, and using the category corresponding to the maximum probability as the image recognition result, which is then output. The classification layer includes a multilayer perceptron network and a Softmax function. The multilayer perceptron network typically consists of multiple fully connected layers and activation functions, used to map high-dimensional feature vectors to a feature space equal to the number of categories. In this invention, it performs a nonlinear transformation on the final token features to obtain a classification feature vector. The Softmax function performs probability mapping on the classification feature vector to obtain the predicted probability distribution for each category.

[0093] Specifically, starting with the characteristics of the final token Extract the feature vector corresponding to the classification token. The feature vector is located at the first position of the sequence and has a dimension of 1×D. This classification token feature vector is then input into a multilayer perceptron network for sublinear transformation to achieve classification. Taking a multilayer perceptron with one hidden layer as an example, the processing procedure is as follows: Classify Token Feature Vectors The input is linearly transformed into the first fully connected layer to obtain the hidden layer features V, as shown in the formula:

[0094] in, Here is the weight matrix of the hidden layer, with dimension D× , This represents the number of hidden layer nodes. This is the bias term of the hidden layer, with a dimension of 1× GELU is the Gaussian error linear unit activation function, used to introduce nonlinear processing.

[0095] Then, the hidden layer features V are input into the second fully connected layer for linear transformation to obtain the classification feature vector f, calculated as follows:

[0096] in, Here is the weight matrix of the output layer, with dimension 1. ×C, where C is the total number of categories in the image classification task; The bias term of the output layer is 1×C, and the dimension of the classification feature vector f is 1×C, with each element corresponding to the original score of a class.

[0097] Next, the obtained classification feature vector f is input into the Softmax function for probability mapping, transforming it into a predicted probability distribution for each category. The Softmax function maps the original score to probability values ​​between 0 and 1, with the sum of the probabilities of all categories being 1. The calculation formula is as follows:

[0098] in, Let f be the j-th element in the classification feature vector f, representing the raw score of the j-th category; exp is the exponential function; C is the total number of categories; Let be the predicted probability of the j-th category.

[0099] The predicted probability distribution vector p = [ is calculated for all categories using the formula above.] , ,..., The dimension of the predicted probability distribution vector is 1×C.

[0100] Furthermore, based on the obtained predicted probability distribution vector p, the category index corresponding to the maximum probability is determined as the image recognition result. The specific operation is as follows: =

[0101] Where argmax represents the index corresponding to the maximum value, and the value of j ranges from 1 to C. This indicates the predicted category number, which corresponds to a category in a predefined set of category labels.

[0102] Finally, the determined image recognition results are output. These results can be category numbers, category names, or other forms of category identifiers. They can be directly used in subsequent application scenarios, such as the display output of image classification systems, the result annotation of object detection systems, and the query input of image retrieval systems.

[0103] In addition, refer to Figure 7 The present invention also provides a pruning device for a Vision Transformer model based on statistical learning. The pruning device is used to implement the pruning method described in the present invention. The pruning device includes: an image token sequence acquisition module 1, a current input token feature sequence acquisition module 2, a current retained token quantity acquisition module 3, a current retained token feature sequence acquisition module 4, and a final token feature acquisition module 5.

[0104] The image token sequence acquisition module 1 is configured to: acquire an initial token sequence based on the original image.

[0105] The current input token feature sequence acquisition module 2 is configured to: acquire the current input token feature sequence based on the current input token sequence through the current encoding layer of the Transformer encoder. It includes a single-head attention output feature acquisition submodule 21, a multi-head attention encoding feature acquisition submodule 22, and an attention encoding feature nonlinear transformation submodule 23. The single-head attention output feature acquisition submodule 21 is used to acquire the single-head attention output feature of each attention head based on the current input token sequence. The multi-head attention encoding feature acquisition submodule 22 is used to obtain the multi-head attention encoding feature based on each single-head attention output feature. The attention encoding feature nonlinear transformation submodule 23 is used to acquire the current input token feature sequence based on the multi-head attention encoding feature.

[0106] The current token retention quantity acquisition module 3 is configured to determine the number of tokens that need to be retained in the current encoding layer based on the preset pruning ratio.

[0107] The current retained token feature sequence acquisition module 4 is configured to: based on the current input token feature sequence and the number of tokens that the current encoding layer needs to retain, obtain the current retained token feature sequence from the current input token feature sequence using a statistical learning method. This module includes a current candidate retained token combination feature acquisition submodule 41, a total loss function value acquisition submodule 42, and a current retained token feature sequence selection submodule 43. Specifically, the current candidate retained token combination feature acquisition submodule 41 generates multiple current candidate retained token combination features based on the current input token feature sequence and the number of tokens that the current encoding layer needs to retain; the total loss function value acquisition submodule 42 acquires the total loss function value corresponding to each candidate retained token combination feature; and the current retained token feature sequence selection submodule 43 selects the candidate retained token combination feature with the smallest total loss function value, thereby obtaining the current retained token feature sequence.

[0108] The final token feature acquisition module 5 is configured to: take the next encoding layer as the current encoding layer, and take the current retained token feature sequence corresponding to the previous encoding layer as the current input token sequence, traverse all encoding layers of the Transformer model, and thus obtain the current retained token feature sequence corresponding to the last encoding layer, thereby obtaining the final token feature.

[0109] In addition, the present invention also provides an image recognition device, which includes the aforementioned pruning device and an image recognition module, wherein the image recognition module is configured to: perform image recognition based on the final token feature and output the image recognition result.

[0110] Compared to existing technologies, the solution of this invention, based on statistical learning, obtains the total loss function value corresponding to the feature combination of each candidate retained token, selects the minimum total loss function value, and then obtains the current retained token feature sequence based on the candidate retained token feature combination with the minimum total loss function value. This reduces the possibility of erroneously deleting important tokens, thereby improving model performance and enhancing the accuracy of results in tasks such as image classification, object detection, and image segmentation. Furthermore, by constructing a total loss function that includes feature reconstruction error terms and regularization terms, it comprehensively evaluates the importance of tokens from two dimensions: information retention and quantity control. This makes pruning decisions more accurate and effectively avoids the problem of erroneously deleting important tokens due to a single evaluation metric, thus maximizing model performance while reducing computational complexity. In addition, by pruning tokens layer by layer, the number of tokens that need to be processed in the Transformer encoder layer is effectively reduced, significantly lowering the model's computational complexity and memory usage. While maintaining high accuracy in classification and detection results, it improves the model's inference speed, making it better suited for resource-constrained mobile devices or applications with high real-time requirements.

[0111] The embodiments described above are merely examples of several implementations of the present invention, and while the descriptions are specific and detailed, they should not be construed as limiting the scope of the invention. It should be noted that those skilled in the art can make various modifications and improvements without departing from the present invention, and the present invention is intended to include these modifications and variations.

Claims

1. A pruning method for ViT models based on statistical learning, characterized in that, Including the following steps: Based on the original image, obtain the initial token sequence; The initial token sequence is used as the current input token sequence, and the current input token feature sequence is obtained through the current encoding layer of the Transformer encoder; Based on the preset pruning ratio, determine the number of tokens that need to be retained in the current encoding layer; Based on the current input token feature sequence and the number of tokens that the current encoding layer needs to retain, the current retained token feature sequence is obtained from the current input token feature sequence using a statistical learning method. The next encoding layer is used as the current encoding layer, and the current retained token feature sequence corresponding to the previous encoding layer is used as the current input token sequence. Steps S2-S4 are repeated to traverse all encoding layers of the Transformer model, thereby obtaining the final token features. The step of obtaining the currently retained token feature sequence from the current input token feature sequence based on statistical learning methods includes: Based on the current input Token feature sequence and the number of Tokens that the current encoding layer needs to retain, generate multiple current candidate retained Token combination features; Obtain the total loss function value corresponding to each current candidate retained token combination feature; Select the current candidate retained token combination feature with the minimum total loss function value to obtain the current retained token feature sequence.

2. The pruning method according to claim 1, characterized in that, The total loss function value is the sum of the feature reconstruction error and the regularization term.

3. The pruning method according to claim 2, characterized in that, The feature reconstruction error is obtained as follows: based on the attention mechanism, the current candidate retained token combination features are aligned to the dimension of the current input token feature sequence to obtain the aligned current candidate retained token combination features; then, the feature reconstruction error between the current input token feature sequence and the aligned current candidate retained token combination features is calculated.

4. The pruning method according to claim 3, characterized in that, The aligned current candidate retained token combination features are obtained by the following steps: calculate the attention weight matrix of the current candidate retained token combination features to the current input token feature sequence, and then generate the aligned current candidate retained token combination features by weighted summation.

5. The pruning method according to claim 4, characterized in that, The feature reconstruction error between the current input token feature sequence and the aligned current candidate retained token combination features is calculated using L2 norm, L1 norm, or cosine similarity.

6. The pruning method according to any one of claims 2 to 5, characterized in that, Calculate the regularization term based on the number of tokens that need to be retained in the current encoding layer.

7. The pruning method according to any one of claims 1 to 6, characterized in that, Obtain the feature sequence of the current input token, including: Based on the current input Token sequence, obtain the single-head attention output features of each attention head; Based on the output features of each single-head attention, the multi-head attention encoding features are obtained; Based on the multi-head attention encoding features, the current input token feature sequence is obtained through a feedforward neural network.

8. The pruning method according to claim 7, characterized in that, The initial token sequence includes a classification token and multiple image patch embedding vectors, and the current candidate retained token combination features include the classification token and multiple image patch token features.

9. A ViT model pruning device based on statistical learning, characterized in that, This device is used to implement the pruning method according to any one of claims 1-8.

10. An image recognition method, characterized in that, It includes the pruning method of claim 8, and further includes: image recognition based on the final token feature.

Images