Probe token-based image feature extraction network construction and training method

Abstract

The invention specifically discloses an image feature extraction network construction and training method based on a probe token. In order to solve the technical problems that a traditional attention mechanism in an existing image feature extraction network is high in calculation complexity and high-resolution images are difficult to process efficiently, probe tokens, the number of which is much smaller than that of input feature tokens, are introduced as intermediaries; attention calculation is divided into two steps of probe aggregation attention and probe broadcast attention. Firstly, global information of input features is aggregated by a probe token, and then the aggregated information is broadcasted to each query token. The method reduces computational complexity to linear complexity while maintaining global context modeling capability. The method can directly replace a standard attention module in an existing visual Transform, is suitable for various visual tasks such as image classification, target detection and semantic segmentation, remarkably improves the calculation efficiency on the premise of ensuring the expression ability of the model, and is particularly suitable for a high-resolution image processing scene.

Description

Technical Field

[0001] This invention belongs to the field of computer vision and deep learning technology, specifically relating to a method for constructing and training an image feature extraction network based on probe tokens, which can be applied to visual tasks such as image classification, object detection, semantic segmentation and image generation. Background Technology

[0002] Feature extraction models, through self-attention mechanisms, can effectively model global dependencies and have achieved remarkable results in the field of computer vision. However, the computational complexity of the traditional Softmax attention mechanism is proportional to the square of the number of input tokens (O(N²)), and the computational cost increases sharply when processing high-resolution images, severely limiting its deployment in practical applications.

[0003] Existing methods for reducing computational complexity mainly include window attention mechanisms that limit the receptive field and sparse attention mechanisms that reduce the number of key values, but these methods sacrifice the model's global modeling ability. While linear attention mechanisms can reduce complexity to linear, the simple mapping function leads to a significant decrease in the model's expressive power.

[0004] In summary, while existing methods for reducing computational complexity in image feature extraction networks have overcome the computational complexity problem to some extent, they sacrifice global modeling capabilities or feature extraction accuracy, making them unsuitable for handling complex real-world scenarios. Therefore, there is an urgent need for a novel image feature extraction network construction method that maintains global modeling capabilities while achieving linear computational complexity. This invention aims to propose a probe-token-based method for constructing and training image feature extraction networks, which can achieve high-precision image feature extraction with linear computational complexity, thus filling the gaps in existing technologies. Summary of the Invention

[0005] In view of this, the technical problem to be solved by the present invention is to provide a method for constructing and training an image feature extraction network based on probe tokens, which achieves efficient global context modeling by introducing probe tokens into the network encoder as a medium for information aggregation and broadcasting.

[0006] The technical solution adopted in this invention is as follows:

[0007] A method for constructing and training an image feature extraction network based on probe tokens includes the following steps:

[0008] Step 1: Obtain the training dataset for the visual task. Collect or prepare image data for the specific visual task and divide it into training set and validation set.

[0009] Step 2: Preprocess the input image data, including image normalization, resizing, block embedding, and position encoding operations, to convert the input image into a feature token sequence;

[0010] Step 3: Construct a visual feature extraction network based on probe tokens. The network structure mainly includes: a feature embedding module, a probe attention-based encoder module, and a task-specific prediction head module. The probe attention-based encoder module introduces probe tokens as information intermediaries and adopts a two-stage attention mechanism to reduce computational complexity.

[0011] Step 4: Train the visual feature extraction network. Input the preprocessed feature token sequence into the network, perform initial projection through the feature embedding module, then perform feature transformation through the probe attention-based encoder module, and finally generate the task output through the prediction head module. Select the corresponding loss function according to the task type to calculate the error between the network output and the ground truth, and update the network parameters using the backpropagation algorithm.

[0012] Step 5: Repeat step 4 until the network loss function value converges to below the preset threshold, or the number of iterations reaches the preset upper limit; save the trained network model parameters.

[0013] Step 6: Load the saved network model parameters, take a new image to be processed, input it into the network after the same preprocessing as in Step 2, and output the visual task processing results.

[0014] Preferably, the encoder module based on probe attention in step three includes a probe aggregation unit and a probe broadcasting unit, and adopts a two-stage attention mechanism: the first stage uses probe tokens to aggregate global information, and the second stage broadcasts the aggregated information to each query location.

[0015] Preferably, the probe tokens are obtained through learnable parameters, and their number is much smaller than the number of input tokens, reducing the computational complexity from quadratic to linear.

[0016] Preferably, the probe attention calculation incorporates a probe bias term obtained from learnable parameters to encode spatial location information. Beneficial effects

[0017] The present invention has the following beneficial effects:

[0018] High computational efficiency: By introducing a small number of probe tokens, the quadratic complexity of attention computation is reduced to linear, significantly reducing computational costs and making it suitable for high-resolution image processing.

[0019] Strong global modeling capability: By maintaining the global receptive field through probe tokens, it can effectively capture long-distance dependencies and avoid information loss.

[0020] Flexible deployment: It can be directly embedded into existing visual feature extraction architectures without requiring significant modifications to the network structure. Attached Figure Description

[0021] Figure 1 Flowchart for network training

[0022] Figure 2 This is a schematic diagram of the overall architecture of a probe token-based visual feature extraction network.

[0023] Figure 3 This is a schematic diagram of the probe attention module, illustrating the two-stage computation process of probe aggregation and probe broadcasting; Detailed Implementation

[0024] To make the objectives, technical solutions, and advantages of the present invention clearer, the specific implementation methods of the present invention will be described in detail below with reference to the embodiments: This embodiment is implemented under the premise of the technical solution of the present invention, and provides detailed implementation methods and specific operation processes, but the protection scope of the present invention is not limited to the following embodiments.

[0025] Example 1: Overall Network Construction and Training

[0026] Step 1: Data Preparation and Preprocessing

[0027] For image classification tasks, prepare the ImageNet dataset or other domain-specific datasets, dividing them into training and validation sets. Preprocess the input images: First, uniformly adjust the image size to 224×224 pixels, and use data augmentation techniques such as random cropping, scaling, and color dithering to improve the network's generalization ability; then, normalize the pixel values ​​of the augmented images, mapping pixel intensity values ​​to the [0,1] interval to unify the numerical range of the input images; then, divide the image into 56×56 image blocks using patch embedding, and convert each block into a feature vector through linear projection; finally, add sinusoidal positional encoding to generate a feature token sequence.

[0028] Step 2: Network Architecture Setup

[0029] A visual feature extraction network based on probe tokens is constructed. The visual feature extraction network mainly includes:

[0030] Feature embedding module: It uses multiple convolutional layers, batch normalization layers, and activation functions to extract local features of the image and projects the input image into a feature sequence with a target embedding dimension of C.

[0031] Specifically, the feature embedding module includes four cascaded convolutional processing sub-layers. Each sub-layer consists of a convolutional layer and a batch normalization layer. The sub-layers are connected via the GELU activation function, forming a progressive feature extraction and downsampling path. The first and third layers are configured with a 3×3 kernel size, a horizontal and vertical stride of 2, 1 pixel edge padding, and output channels of [number missing]. and With this configuration, the spatial resolution of the feature map is reduced to half that of the input image. The second and fourth layers are configured with: a 3×3 kernel size, a horizontal and vertical stride of 1, 1 pixel edge padding, and output channels of [missing information]. and With this configuration, the spatial resolution of the feature map remains unchanged.

[0032] Preferably, the sub-layers at each level are connected by the GELU activation function, the mathematical expression of which is: ,in The cumulative distribution function is the standard normal distribution. Compared with the ReLU activation function, it has smoother gradient characteristics, which is beneficial to the training stability of deep networks.

[0033] The encoder module based on probe attention: The encoder module based on probe attention adopts a two-stage attention mechanism by using a probe aggregation attention subnetwork and a probe broadcast attention subnetwork.

[0034] The first phase is probe aggregation attention: the probe aggregation attention subnetwork receives probe tokens. and input features As input, the probe token A is used to generate a query matrix through linear projection. , and All dimensions are The input features I are used to generate key matrices through linear projection. Sum matrix ,in , , All dimensions are The query matrix generated by the probe token. As a query, with the key matrix Sum matrix Perform probe aggregation attention calculation to generate aggregated probe features. ;

[0035] The probe aggregation attention calculation is expressed as follows:

[0036]

[0037] Where Softmax is the normalization exponential function. · denotes matrix transpose, and · denotes matrix multiplication;

[0038] The second stage is probe broadcast attention: the probe broadcast attention subnetwork receives aggregated probe features. and input features As input, the input feature I is used to generate a query matrix through linear projection. , The dimension of w is N×C. The aggregated probe features... The bond matrices are generated by linear projection. Sum matrix ,in , The dimensions of all are n×C; the input features are... The generated query matrix As a query, with the key matrix Sum matrix Perform Galerkin broadcast attention calculation with Softmax to generate output image features. ;

[0039] The Galerkin-type broadcast attention calculation with Softmax is represented as follows:

[0040]

[0041] in, For normalized exponential functions, Indicates matrix transpose. Represents matrix multiplication. For layer normalization operation;

[0042] Feature recovery module: The value matrix is ​​processed by a depthwise separable convolution module and added to the output of the probe attention computation sublayer to preserve feature diversity and detail information.

[0043] Prediction head module: Selected according to the task type. For image classification tasks, global average pooling followed by fully connected layers is used; for dense prediction tasks, deconvolution or upsampling modules are used.

[0044] Step 3: Model Training

[0045] The classification error was calculated using the cross-entropy loss function. The AdamW optimizer was used with an initial learning rate of 1×10^-3, employing a cosine annealing learning rate scheduling strategy, and training was performed for 300 epochs. Stochastic depth regularization was used to prevent overfitting.

[0046] Step 4: Model Saving and Deployment

[0047] When the validation set accuracy no longer improves or reaches the preset training epochs, the model parameters are saved. During the inference phase, pre-trained weights are loaded, and the same preprocessing as during training is performed on new input images. The prediction results are then output through forward computation of the network.

[0048] The specific embodiments described above further illustrate the purpose, technical solution, and beneficial effects of the present invention. It should be understood that the above descriptions are merely specific embodiments of the present invention and are not intended to limit the invention. Any modifications, equivalent substitutions, or improvements made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

Claims

1. A method for constructing and training an image feature extraction network based on probe tokens, characterized in that, Includes the following steps: Step 1: Obtain the visual task training dataset. Collect or prepare image data for the specific visual task and divide it into training and validation sets. Step 2: Preprocess the input image data, including image normalization, resizing, block embedding and position encoding operations, to convert the input image into a feature token sequence. Step 3: Construct a probe token-based visual feature extraction network. The network structure mainly includes a feature embedding module, a probe attention-based encoder module, and a task-specific prediction head module. The probe attention-based encoder module introduces probe tokens as information intermediaries and adopts a two-stage attention mechanism to reduce computational complexity. Step 4: Train the visual feature extraction network. The specific method is as follows: Input the preprocessed feature token sequence into the network, perform initial projection through the feature embedding module, then perform feature transformation through the probe attention-based encoder module, and finally generate the task output through the prediction head module. Select the corresponding loss function according to the task type to calculate the error between the network output and the ground truth, and update the network parameters using the backpropagation algorithm. Step 5: Repeat Step 4 until the network loss function value converges to below a preset threshold, or the number of iterations reaches a preset upper limit. Save the trained network model parameters. Step 6: Load the saved network model parameters, take a new image to be processed, perform the same preprocessing as in Step 2, input it into the network, and output image features for downstream task network processing.

2. The image feature extraction and network training method based on probe tokens according to claim 1, characterized in that, The probe-attention-based encoder module described in step three includes the following structure: Probe Aggregation Attention Subnetwork: The probe aggregation attention subnetwork receives probe tokens and input features and generates a query matrix, a key matrix, and a value matrix; The global information of the feature tokens is aggregated into the probe token using Softmax attention to obtain the aggregated probe token; Probe Broadcast Attention Subnetwork: The probe broadcast attention subnetwork receives aggregated probe tokens and input features and generates query matrix, key matrix and value matrix; it uses Galerkin attention with Softmax to distribute the information of the aggregated probe tokens to each input feature query position to obtain output image features; Feature recovery subnetwork: used to maintain feature diversity and fuse multi-scale information, outputting the transformed feature representation.

3. The method according to claim 2, characterized in that, The encoder module for probe attention employs a probe aggregation attention subnetwork and a probe broadcast attention subnetwork to implement a two-stage attention mechanism: The first stage is probe aggregation attention: the probe aggregation attention subnetwork receives probe token A and input feature I as input. The probe token A is then used to generate a query matrix through linear projection. A and The dimensions of each feature are n×C; the input features I are used to generate key matrices through linear projection. Sum matrix Among them, I, , All dimensions are N×C. The query matrix generated from the probe token... As a query, with the key matrix Sum matrix Perform probe aggregation attention calculation to generate aggregated probe features. ; The probe aggregation attention calculation is expressed as follows: ， Where Softmax is the normalization exponential function. · denotes matrix transpose, and · denotes matrix multiplication; The second stage is probe broadcast attention: the probe broadcast attention subnetwork receives aggregated probe features. And input feature I is used as input. The input feature I is then used to generate a query matrix through linear projection. , The dimension of w is The polymer probe features The bond matrices are generated by linear projection. Sum matrix ,in , All dimensions are The query matrix generated from input feature I As a query, with the key matrix Sum matrix Perform Galerkin broadcast attention calculation with Softmax to generate output image features. ; The Galerkin-type broadcast attention calculation with Softmax is represented as follows: ， in, For normalized exponential functions, ∠ represents matrix transpose, · represents matrix multiplication, and Ln is the layer normalization operation.

4. The method according to claim 3, characterized in that, The probe tokens are initialized using learnable parameters; the number of probe tokens n is much smaller than the number of input feature tokens N, reducing computational complexity from... Reduce to Set n to 50 if the input image resolution is 224×224 pixels.

5. The method according to claim 3, characterized in that, The probe aggregation attention calculation and probe broadcast attention calculation respectively introduce probe bias terms to encode spatial location information; the probe bias is obtained through learnable parameters, and its dimension is the same as that of the probe token.

6. The method according to claim 2, characterized in that, The feature recovery subnetwork uses a depthwise separable convolutional module to process the value matrix and adds it to the output of the probe attention computation subnetwork to maintain feature diversity and detailed information.

Images