An image classification method based on probability graph and ViT model

Abstract

This invention relates to an image classification method based on a probabilistic graph and a ViT model, belonging to the technical field of computer vision image classification. This method models a multi-head attention mechanism from the perspective of probability theory, treating the attention values in the multi-head attention as latent variables. By leveraging the Explaining-Away Effects of the probabilistic graph model and the hierarchical structure of the Transformer, attention logits are transferred layer by layer and the values of adjacent layers are fused, promoting interaction between different heads. This invention addresses the problem of head parameter redundancy in the multi-head attention mechanism in the conventional Vision Transformer model. By modeling the multi-head attention mechanism as a probabilistic graph model, treating the attention values as latent variables, and promoting interaction between different attention heads.

Description

Technical Field

[0001] The present invention relates to an image classification method based on a probability graph and a ViT model, and belongs to the technical field of computer vision image classification. Background Art

[0002] The Transformer's tremendous success in natural language processing has inspired attempts to introduce multi-head attention mechanisms into computer vision, where convolutional neural networks are the dominant framework. Compared to convolutional neural networks, the Transformer offers significant advantages in capturing global information in images. Furthermore, its parallelizable computation has facilitated its application in vision. Currently, the Vision Transformer has achieved remarkable results in various computer vision tasks, such as image classification, object detection, and image denoising. However, numerous studies have found that parameter redundancy between the different heads in the Transformer's core multi-head attention mechanism significantly impacts the model's overall performance. Summary of the Invention

[0003] The purpose of the present invention is to provide an image classification method based on probability graph and ViT model to solve the above-mentioned defects.

[0004] In order to achieve the above object, the technical solution adopted by the present invention is:

[0005] A method for image classification based on probability graph and ViT model includes the following steps:

[0006] S1. Divide the image input to the model into blocks, then flatten each image block into a one-dimensional vector, and finally generate a patch embedding through linear transformation;

[0007] S2, add position encoding to each patch embedding to supplement the position information;

[0008] S3: Add a token for classification and learn the overall information of other image patches;

[0009] S4, the Transformer Block based on head interaction, treats attention values as latent variables, uses the Explaining-Away Effects in the probabilistic graphical model and the hierarchical structure of the Transformer to pass attention logits layer by layer and fuse the values of adjacent layers to promote interaction between different heads;

[0010] S5. Use two fully connected layers to input the classification token into the classification layer to obtain the classification result of the image.

[0011] A further improvement of the technical solution of the present invention is that the specific steps of S1 are:

[0012] S11. Divide and flatten the image input to the model. The specific operations are as follows:

[0013] Set the length and width of the image patch to P, that is, transform the image data H*W*C into

[0014]

[0015] Where N is the number of patches for an image segmentation, C is the number of channels, H is the image height, and W is the image width;

[0016] S12. Linearly transform the patch vector into a patch embedding:

[0017] patch_embedding=nn.Linear(patch_dim,dim)

[0018] Among them, patch_dim is the dimension of the patch vector, and dim is the dimension of the patch embedding.

[0019] A further improvement of the technical solution of the present invention is that the specific operation of S2 is:

[0020] pos_embedding=nn.Parameter(torch.randn(1,num_patches+1,dim))

[0021] Among them, pos_embedding is the position encoding of the patch, and num_patches is the number of patches.

[0022] A further improvement of the technical solution of the present invention is that the specific operation of S3 is:

[0023] Add a learnable encoding specifically for classification and concatenate it with the input, specifically:

[0024] cls_token=nn.Parameter(torch.randn(1,1,dim))

[0025] Among them, cls_token is the classification token, which is then spliced with other patch tokens.

[0026] A further improvement of the technical solution of the present invention is that the specific steps of S4 are:

[0027] S41, attention head sequence modeling, regards attention value as a latent variable,

[0028] p(Y|X)=∫ A p(Y|A,X)p(A|X)dA

[0029] Where Y is the image label, X is the input photo, A is the middle layer attention value, and p(A|X) is the joint prior distribution;

[0030] S42, Transformer hierarchical modeling, using the hierarchical structure of transformer, this process can be expressed as:

[0031]

[0032] Among them, A j represents the attention distribution of the jth layer, A j The calculation depends on A j-1 The recursive structure, specifically, A j The calculation formula is as follows:

[0033]

[0034] S43, attention fusion of adjacent layers, adding MLP to the multi-head attention calculation module of the Transformer layer block, fusing the attention vlaues between layers and promoting the de-redundancy of different heads. This process can be expressed as:

[0035] A j =softmax(z j +MLP(z j ,z j-1 ))

[0036] Among them, z j is the attention logits of the jth layer, and MLP is a two-layer fully connected layer used to fuse the attention values of adjacent layers.

[0037] A further improvement of the technical solution of the present invention is that the specific operation formula of S5 is:

[0038] x = self.to_cls_token(x[:,0])

[0039] y = self.mlp_head(x)

[0040] Among them, x is the output classification Token, mlp_head() is the classification layer, and y is the output prediction.

[0041] Due to the adoption of the above technical solution, the technical effects achieved by the present invention are as follows:

[0042] The present invention designs an image classification method based on probability graph and ViT model to address the redundancy problem of head parameters of the multi-head attention mechanism in the ordinary Vision Transformer model. The multi-head attention mechanism is modeled as a probability graph model, and the attention value is regarded as a latent variable to promote the interaction between different attention heads.

[0043] In order to promote the interaction between different heads, the present invention transfers the attention logits layer by layer and fuses the attention logits between adjacent layers, so that different heads can capture different features.

[0044] The present invention improves parameter efficiency by promoting head interaction, thereby improving the accuracy of graphic classification and the effect of transfer learning, while also improving the interpretability of features. BRIEF DESCRIPTION OF THE DRAWINGS

[0045] Figure 1 It is the algorithm flow chart of the present invention;

[0046] Figure 2 It is a model architecture diagram of the present invention; DETAILED DESCRIPTION

[0047] The present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments:

[0048] A method for image classification based on probability graph and ViT model, such as Figure 1 As shown, the following steps are included:

[0049] S1, image segmentation;

[0050] The input image of the model is divided into blocks, each image block is flattened into a one-dimensional vector, and finally a linear transformation is performed to generate a patch embedding;

[0051] S2, position coding;

[0052] Add position encoding to each patch embedding to supplement the position information;

[0053] S3, classification token;

[0054] Add a token for classification to learn the overall information of other image patches;

[0055] S4, Transformer Block based on head interaction;

[0056] Treating attention values as latent variables, we leverage Explaining-Away Effects in probabilistic graphical models and the hierarchical structure of Transformers to propagate attention logits layer by layer and fuse values from adjacent layers to promote interaction between different heads.

[0057] S5. Use two fully connected layers to input the classification token into the classification layer to obtain the classification result of the image.

[0058] The specific implementation process is as follows:

[0059] S1. Image segmentation

[0060] S11. Divide and flatten the image input to the model. This process can be described as:

[0061] Set the length and width of the image patch to P, that is, transform the image data H*W*C into

[0062]

[0063] Among them, N is the number of patches of an image segmentation, C is the number of channels, H is the image height, and W is the image width.

[0064] S12, linearly transform the patch vector into patch embedding. This process can be described as:

[0065] patch_embedding=nn.Linear(patch_dim,dim)

[0066] Among them, patch_dim is the dimension of the patch vector, and dim is the dimension of the patch embedding.

[0067] S2, position encoding, to preserve the relative position information of the image, add a learnable position encoding to each patch

[0068] pos_embedding=nn.Parameter(torch.randn(1,num_patches+1,dim))

[0069] Among them, pos_embedding is the position encoding of the patch, and num_patches is the number of patches.

[0070] S3, classification token, adds a learnable code specifically for classification and concatenates it with the input. Specifically:

[0071] cls_token=nn.Parameter(torch.randn(1,1,dim))

[0072] Among them, cls_token is the classification token, which is then spliced with other patch tokens.

[0073] S4, Transformer Block based on head interaction

[0074] S41, attention head sequence modeling, regards attention value as a latent variable,

[0075] p(Y|X)=∫ A p(Y|A,X)p(A|X)dA

[0076] Among them, Y is the image label, X is the input photo, A is the middle layer attention value, and p(A|X) is the joint prior distribution.

[0077] S42, Transformer hierarchical modeling, using the hierarchical structure of transformer, this process can be expressed as:

[0078]

[0079] Among them, A j represents the attention distribution of the jth layer, A j The calculation depends on A j-1 The recursive structure, specifically, A j The calculation formula is as follows:

[0080]

[0081] S43, attention fusion of adjacent layers, adding MLP to the multi-head attention calculation module of the Transformer layer block, fusing the attention vlaues between layers and promoting the de-redundancy of different heads. This process can be expressed as:

[0082] A j =softmax(z j +MLP(z j ,z j-1 ))

[0083] Among them, z jis the attention logits of the jth layer, and MLP is a two-layer fully connected layer used to fuse the attention values of adjacent layers.

[0084] S5. Use two fully connected layers to input the classification token into the classification layer to obtain the classification result of the image:

[0085] x = self.to_cls_token(x[:,0])

[0086] y = self.mlp_head(x)

[0087] Among them, x is the output classification Token, mlp_head() is the classification layer, and y is the output prediction.

[0088] The present invention proposes to model the multi-head attention mechanism from the perspective of probability theory. The attention value in the multi-head attention is regarded as a latent variable. The Explaining-away Effects of the probabilistic graphical model and the hierarchical structure of the Transformer are used to pass the attention logits layer by layer, and the values of adjacent layers are fused to promote interaction between different heads. In the experiment, the model designed by the method of the present invention was first pre-trained on the ImageNet1K dataset, and the image classification accuracy was improved by 1.56% compared with the ordinary ViT model; then the trained model parameters were applied to the downstream small dataset CIFAR100 dataset for transfer learning, and the classification accuracy was improved by 3.47%. Experiments have shown that the method of the present invention has promoted the effectiveness of parameters and the independence between different heads to a certain extent, and improved the accuracy of image classification and transfer learning performance. On the other hand, by visualizing the attention of the image, it was found that the method of the present invention enables different heads to pay attention to different parts of the image, proving that model learning has certain research significance and application value in terms of feature interpretability.

Claims

1. A method for image classification based on probability graph and ViT model, characterized in that: The following steps are involved: S1. Divide the image input to the model into blocks, then flatten each image block into a one-dimensional vector, and finally generate a patch embedding through linear transformation; S2, add position encoding to each patch embedding to supplement the position information; S3: Add a token for classification and learn the overall information of other image patches; S4, the Transformer Block based on head interaction, treats attention values as latent variables, uses the Explaining-Away Effects in the probabilistic graphical model and the hierarchical structure of the Transformer to pass attention logits layer by layer and fuse the values of adjacent layers to promote interaction between different heads. The specific steps of S4 are: S41, attention head sequence modeling, regards attention value as a latent variable, p(Y∣X)=∫ A p(Y∣A,X)p(A∣X)dA Where Y is the image label, X is the input photo, A is the middle layer attention value, and p(A|X) is the joint prior distribution; S42, Transformer hierarchical modeling, using the hierarchical structure of transformer, this process can be expressed as: Among them, A j represents the attention distribution of the jth layer, A j The calculation depends on A j-1 The recursive structure, specifically, A j The calculation formula is as follows: S43, attention fusion of adjacent layers, adding MLP to the multi-head attention calculation module of the Transformer layer block, fusing the attention vlaues between layers and promoting the de-redundancy of different heads. This process can be expressed as: AND j =softmax(z j +MLP(z j ,With j-1 )) Among them, z j is the attention logits of the jth layer, and MLP is a two-layer fully connected layer used to fuse and interact the attention values of adjacent layers; S5. Use two fully connected layers to input the classification token into the classification layer to obtain the classification result of the image.

2. The image classification method based on probability graph and ViT model according to claim 1, characterized in that: The specific steps of S1 are: S11. Divide and flatten the image input to the model. The specific operations are as follows: Set the length and width of the image patch to P, that is, transform the image data H*W*C into Where N is the number of patches for an image segmentation, C is the number of channels, H is the image height, and W is the image width; S12. Linearly transform the patch vector into a patch embedding: patch_embedding=nn.Linear(patch_dim,dim) Among them, patch_dim is the dimension of the patch vector, and dim is the dimension of the patch embedding.

3. The image classification method based on probability graph and ViT model according to claim 1, characterized in that: The specific operations of S2 are: pos_embedding = nn.Parameter(torch.randn(1,num_patches+1,dim)) where pos_embedding is the position encoding of the patch and num_patches is the number of patches.

4. The image classification method based on probability graph and ViT model according to claim 1, characterized in that: The specific operations of S3 are: Add a learnable encoding specifically for classification and concatenate it with the input, specifically: cls_token=nn.Parameter(torch.randn(1,1,dim)) Among them, cls_token is the classification token, which is then spliced with other patch tokens.

5. The image classification method based on probability graph and ViT model according to claim 1, characterized in that: The specific operation formula of S5 is: x = self.to_cls_token(x[:,0]) y = self.mlp_head(x) Among them, x is the output classification Token, mlp_head() is the classification layer, and y is the output prediction.

Images