A method for fine-tuning a visual Mamba model based on selective cues

Abstract

This invention discloses a method for efficient fine-tuning of a visual Mamba model based on selective prompts, belonging to the field of image processing technology. This invention transforms the fine-tuning process of a pre-trained visual Mamba model into a prompt-based information propagation optimization method by introducing a selective prompt module. This method utilizes the selective prompt module to generate adaptive prompts based on the sample dataset, enhancing the model's responsiveness to input data and thus effectively improving the performance of classification tasks. Specifically, this invention employs a dual-path structure, including cross-layer prompts and intra-layer prompts, optimizing the propagation of shared information between model layers and specific information within layers, respectively. Cross-layer prompts achieve information sharing between different layers, ensuring feature consistency; secondly, intra-layer prompts focus on extracting and maintaining specific features of each layer, thereby improving the model's discriminative ability. This method not only avoids the problem of catastrophic forgetting but also improves the training efficiency and performance of the visual Mamba model in downstream classification tasks.

Description

Technical Field

[0001] This invention belongs to the field of image processing technology, specifically relating to a method for fine-tuning a visual Mamba model based on selective cues. Background Technology

[0002] With the widespread application of large-scale pre-trained models, model fine-tuning in visual tasks has become one of the important techniques in the field of computer vision. In tasks such as image classification and object detection, pre-trained models provide excellent feature representation capabilities, which can effectively improve the performance of downstream tasks. Therefore, how to efficiently fine-tune pre-trained models has become a hot research topic.

[0003] Existing methods for fine-tuning visual models are typically based on models such as Vision Transformer (ViT) or Convolutional Neural Networks (CNN), relying on adjusting the entire model or some parameters to adapt to new downstream tasks. However, as the model size increases, directly fine-tuning all parameters not only requires a large amount of computational resources but also leads to the problem of catastrophic forgetting, where the model loses its memory of previous tasks when learning new tasks, affecting the performance of downstream tasks.

[0004] To address this issue, the Parameter Efficient Fine-Tuning (PEFT) method has emerged. Visual Prompting, a technique within this method, introduces a small number of learnable parameters into the pre-trained model, achieving efficient model adaptation while conserving computational resources. However, most existing visual prompting methods are based on ViT models and do not fully consider the characteristics of visual Mamba models in sequence compression and propagation, resulting in poor performance in applications using visual Mamba models. Therefore, a fine-tuning method specifically for visual Mamba models is urgently needed to improve their information propagation capabilities and performance in downstream vision tasks. Summary of the Invention

[0005] To improve the fine-tuning efficiency and performance of large-scale pre-trained visual Mamba models in downstream image classification tasks, this invention proposes a method for fine-tuning visual Mamba models based on selective prompts. This method not only generates prompts adapted to the input data but also optimizes the information propagation path by selectively activating the model's update and forget gates, effectively improving the performance of classification tasks and achieving efficient model fine-tuning while avoiding catastrophic forgetting.

[0006] A method for fine-tuning a visual Mamba model based on selective cues includes:

[0007] Obtain a sample dataset, which includes sample images and corresponding classification labels for the sample images;

[0008] Selective cueing modules are added to each layer of the visual Mamba model. The modules include cross-layer cueing generators and intra-layer cueing generators. The cross-layer cueing generators share parameters between layers and generate cross-layer cueing information. The intra-layer cueing generators extract and retain features of each layer and generate intra-layer cueing information.

[0009] The sample image is converted into an image token and input into the visual Mamba model along with a randomly initialized category token. The image token is propagated through each layer, with the cross-layer prompts and the intra-layer prompts superimposed, until the last layer of the model. The category token is propagated through each layer until the last layer of the model.

[0010] The output category token of the last layer of the model is passed to the classification head to generate the probability distribution of classification prediction;

[0011] Calculate the loss function based on the probability distribution and the classification label;

[0012] By fixing the pre-training parameters of the visual Mamba model, the cross-layer cue generator, the intra-layer cue generator, and the classification head are trained based on the loss function and the sample dataset until convergence, resulting in a visual Mamba model with selective cueing.

[0013] Furthermore, the cross-layer prompt generator and the intra-layer prompt generator each include a dynamic adjustment factor that scales their respective prompt information, and the scaled prompt information is superimposed on the input of the current layer.

[0014] Furthermore, each of the in-layer prompt generators includes a linear dimensionality reduction layer, a linear dimensionality increase layer, and a SiLU activation function.

[0015] Furthermore, each of the aforementioned cross-layer cue generators includes a fully connected generator and a SiLU activation function.

[0016] Furthermore, the activation status of the update gate is analyzed using visualization tools.

[0017] Furthermore, the sample dataset contains diverse sample images and annotation information.

[0018] Furthermore, the image token is obtained from the sample image by the following method: the sample image is divided into image blocks, and the image token is obtained by block encoding.

[0019] An image classification method that uses the above-mentioned finely tuned selectively prompted visual Mamba model for image classification.

[0020] The beneficial effects of this invention are as follows:

[0021] 1) This invention proposes a method for fine-tuning a visual Mamba model based on selective prompts. By introducing a selective prompt module, it is possible to effectively generate prompt information that adapts to the input data, improve the model's responsiveness to the input data, and thus optimize the performance of classification tasks.

[0022] 2) This invention adopts a dual-path structure, including cross-hints and internal hints, which are optimized for inter-layer shared information and intra-layer specific information, respectively, improving the efficiency of information propagation, ensuring the consistency and distinguishability of features, and avoiding catastrophic forgetting problems.

[0023] 3) By fixing the parameters of the pre-trained model and only fine-tuning the selective prompting module and classification head, this invention greatly reduces the computational and storage resource requirements, while improving the training efficiency and overall performance of the model in downstream tasks. Attached Figure Description

[0024] Figure 1 This is a flowchart of an embodiment of the present invention.

[0025] Figure 2 This is a schematic diagram of the selectively prompted visual Mamba model of the present invention.

[0026] Figure 3 This is a visualization of the information propagation process in the SVMamba model according to an embodiment of the present invention.

[0027] Figure 4 This is a visualization of the update gate in the SVMamba model according to an embodiment of the present invention. Detailed Implementation

[0028] The present invention will now be described in further detail with reference to the accompanying drawings. The examples given are only for explaining the present invention and are not intended to limit the scope of the present invention.

[0029] To enable those skilled in the art to better understand the technical solutions in the embodiments of the present invention, and to make the above-mentioned objectives, features, and advantages of the present invention more apparent and understandable, the technical solutions of the present invention will be further described in detail below with reference to the accompanying drawings:

[0030] The specific implementation flowchart of the method of this invention is as follows: Figure 1 As shown, the process is as follows:

[0031] Step 1) Obtain the dataset, which contains image samples x and corresponding classification labels y for the classification task.

[0032] This invention selected two standard datasets, HTA and VTAB-1K, which are widely used in image classification tasks to train and validate the model's performance. HTA (Head Tuning Adaptation) is a large-scale benchmark dataset containing 10 image classification sub-tasks, each representing a different visual category and scene. This dataset includes CIFAR-10, CIFAR-100, DTD, CUB-200, Stanford-Dogs, Oxford-Flowers, Food101, GTSRB, SVHN, and NABIRDS datasets. The image categories cover natural objects, artificial objects, and signs, and it has a wide range of applications, as detailed in Table 1.

[0033]

[0034] Table 1

[0035] VTAB-1K (Visual Task Adaptation Benchmark) is a visual task adaptation benchmark consisting of 19 tasks divided into three categories: natural (e.g., CIFAR-100), specialized tasks (e.g., EuroSAT), and structured tasks (e.g., Clevr-Count). Each task has 1,000 training samples. These datasets are used to evaluate the model's adaptation performance on diverse downstream tasks.

[0036] Step 2) Add a selective cueing module to the visual Mamba model. Construct a pre-trained model framework that includes the visual Mamba model backbone, and introduce a dual-path selective cueing module; the selective cueing module is used to generate adaptive cueing information to enhance the information dissemination effect, such as... Figure 2 As shown.

[0037] Mamba and the visual Mamba model are based on the Selective State-Space Model (SSM), which offers significant linear computational efficiency and excellent performance, making it particularly suitable for processing long sequence data. The state-space model represents a one-dimensional function or sequence... Mapped to Through N-dimensional hidden state Processing is performed. The hidden state evolves over time, following a linear ordinary differential equation with parameters as follows: and

[0038] h′(t)=Ah(t)+Bx(t)

[0039] y(t)=Ch(t)

[0040] To apply the state-space model to deep learning, zero-order preservation was used for discretization. Continuous parameters A and B are represented by time-scale parameters. Convert to the corresponding discrete form: and

[0041]

[0042] Therefore, the discretized state-space model can be written as:

[0043]

[0044] y i =Ch i

[0045] in Represents an N-dimensional hidden state.

[0046] This invention constructs a pre-trained model framework that includes a visual Mamba model backbone and introduces a dual-path selective cueing module. The core idea of ​​this model is to dynamically adjust the model's processing of input information through an input-dependent modeling mechanism, thereby improving the feature extraction effect.

[0047] During the construction process, a visual Mamba model pre-trained on a large-scale dataset (ImageNet-1K) is first loaded as the backbone. This model not only has deep feature extraction capabilities but also effectively compresses the input sequence, achieving efficient information propagation. Based on the visual Mamba model, this invention introduces a dual-path selective prompting module, which consists of a cross-layer prompting module and an intra-layer prompting module, resulting in a selectively prompted optimized visual Mamba model (SVMamba model).

[0048] The Cross-Prompting Module is designed to capture and propagate shared information between layers to ensure feature consistency. Within this module, the present invention utilizes a fully connected cross-prompting generator for inter-layer shared information. This generator shares parameters across layers, via input x. i Prompt for generating cross-layer information in This represents the input features. The number of layers sharing parameters is a hyperparameter that can be set to 6, 8, or 12. The cross-layer cue generator is a 384-to-384 fully connected generator followed by a SiLU activation function.

[0049] The Inner-Prompting Module focuses on extracting and preserving specific features of each layer to enhance the model's discriminative power. To minimize adjustable parameters while maintaining performance, this invention designs a lightweight Inner-Prompting Generator for each layer of the visual Mamba model. This generator can generate different in-layer hints. The in-layer hint generator contains a linear dimensionality reduction layer (L down ), linear dimensionality increase layer (L up It also includes a SiLU activation function, with the middle hidden dimension set to 64.

[0050] Step 3) Train the model based on the dataset. Based on the sample dataset, the selective prompting module generates prompts to activate the update and forget gates in the visual Mamba model, thereby promoting the propagation of effective information and preventing irrelevant information from interfering with the model's classification process. The cross-entropy loss function is used to calculate the error between the model's predicted category and the actual label.

[0051] Obtain the input image x∈R H×W×C Where (H, W) represents the size of image x, and C is the number of channels. The image is randomly cropped to a size of (224, 224), and then divided into (16, 16) image blocks. After this processing, multiple small blocks are obtained. These small blocks are embedded into a d-dimensional latent space and converted into image tokens through block encoding, represented as...

[0052] The image token is input into the SVMamba model, and propagation proceeds through each layer, combining cross-layer and intra-layer prompts, until the final layer of the optimized model. The selective prompting module in each layer of the model responds to the input x of the current layer. i Generate corresponding cross-layer prompts. and in-layer prompts This can be described by a formula:

[0053]

[0054] Furthermore, considering the different importance of intra-layer and cross-layer prompts, this invention designs two element-wise dynamic scaling factors (α, β) to balance the impact of these prompts on the input distribution.

[0055]

[0056] in and These are learnable parameters, initialized to zero, and the symbol ⊙ represents the Hadamard product.

[0057] Generated prompts The input x is then superimposed onto the current layer. i superior.

[0058]

[0059] In the current layer, the output image token of the selective suggestion module With category token Input together the pre-trained Mamba blocks in the current layer To extract features, the category token is randomly obtained at the beginning of the model. Following the process described above, image tokens and category tokens traverse each layer of the model, incorporating selective cues into the input of each layer. Finally, c is output from the last Mamba block. N+1 Passed to the classification head In the middle, the probability distribution y of the prediction is generated. pred .

[0060] The pre-trained parameters of the visual Mamba model backbone are fixed, and only the selective cueing module and the classification head are trained. Based on the classification loss function and the sample dataset, the model is trained until convergence, resulting in an optimized image classification model.

[0061] This invention calculates the classification loss function based on the probability distribution predicted by the model and the classification label y.

[0062] The selectively prompted visual Mamba model proposed in this invention introduces only a small number of learnable parameters, including the prompt generator in the selective prompting module. And dynamic scaling factors α and β.

[0063]

[0064] During training, this invention keeps the encoder parameters of the pre-trained visual Mamba model unchanged, only allowing the classification head to... and newly added components Training is then performed. Based on the classification loss function and the dataset, the selective cueing module and classification head of the model to be trained are trained until convergence, resulting in a trained SVMamba model.

[0065]

[0066] Step 4) Validate the performance of the trained SVMamba model by testing the model's classification accuracy using benchmark datasets for multiple classification tasks. Simultaneously, analyze the activation status of the update gate using visualization tools to verify the optimization effect of the selective prompting module on information propagation and the improvement of the model's classification performance.

[0067] Finally, this invention performs performance validation on the trained SVMamba model to ensure its effectiveness in downstream classification tasks. This invention uses benchmark datasets for multiple classification tasks to comprehensively test the model's classification accuracy and compares its performance with existing visual cueing methods in image classification tasks, as shown in Table 2. As described in step one, these benchmark datasets include HTA and VTAB-1K, covering a variety of visual tasks and effectively evaluating the model's performance in different scenarios.

[0068] During model evaluation, this invention calculates the classification accuracy of the model on the test set and compares it with existing state-of-the-art models to verify the effectiveness of the proposed method. The formula for calculating classification accuracy is:

[0069]

[0070] Table 2

[0071] Furthermore, this invention uses visualization tools to analyze the activation status of the update gate in the SVMamba model to gain a deeper understanding of the optimization effect of the selective prompting module on information propagation. For example... Figure 3 As shown, this invention first visualizes the changes in the normalized update gate in the SVMamba model during sequence propagation, verifying the effectiveness of this invention in promoting model information propagation. This visualization reveals the information retained in the hidden states of the SVMamba model as the sequence progresses, preventing the loss of discriminative information during state space propagation. The SVMamba model employs a lightweight selective cueing mechanism to generate token-level cueing throughout the sequence, ensuring adaptive activation of the update and forget gates in Mamba blocks to promote the propagation of discriminative information. This highlights the improvements of the method in extracting and transmitting discriminative features.

[0072] Figure 4 Further examples of the update gate in the SVMamba model are shown. The method of this invention is able to retain more discriminative information related to key elements (such as flowers and food), while retaining less discriminative elements such as background (such as tableware) to a lower degree. These visualizations demonstrate that the method of this invention effectively activates the update gate in downstream tasks, enabling the SVMamba model to focus more on discriminative regions, thereby enhancing feature extraction and propagation.

[0073] Finally, through visual analysis of the model's classification performance and information propagation, this invention can comprehensively evaluate the optimization effect of the selective prompting module on improving the model's classification performance. This step provides valuable reference for further model improvement and practical application.

[0074] Although specific embodiments of the invention have been disclosed for illustrative purposes to aid in understanding and implementing the invention, those skilled in the art will understand that various substitutions, variations, and modifications are possible without departing from the spirit and scope of the invention and the appended claims. Therefore, the invention should not be limited to the content disclosed in the preferred embodiments, and the scope of protection claimed by the invention is defined by the claims.

Claims

1. A method for fine-tuning a visual Mamba model based on selective cues, wherein the visual Mamba model consists of N Mamba layers and a classifier head, and each Mamba layer contains... One Mamba block, One Mamba block is used for processing. The method includes: an image token and a Mamba block for processing category tokens. Obtain a sample dataset, which includes sample images and corresponding classification labels for the sample images; Construct a visual Mamba optimization model, which is an Mamba model with an addition before each Mamba layer. A selective prompting module, the selective prompting module including a cross-layer prompting generator and an intra-layer prompting generator; The cross-layer cue generator shares parameters between Mamba layers to generate cross-layer cue information; the intra-layer cue generator extracts and retains features of each Mamba layer to generate intra-layer cue information; the cross-layer cue generator includes a fully connected generator and a SiLU activation function; the intra-layer cue generator includes a linear dimensionality reduction layer, a linear dimensionality increase layer, and a SiLU activation function. Convert the sample image into An image token, along with a randomly initialized category token, is input into the visual Mamba optimization model. The image token propagates through each Mamba layer, superimposed with cross-layer prompts and intra-layer prompts, to the last Mamba layer of the visual Mamba optimization model. The category token propagates through each Mamba layer to the last Mamba layer of the visual Mamba optimization model. The result category token output from the last Mamba layer of the visual Mamba optimization model is passed to the classification head to generate a probability distribution for classification prediction; the loss function is calculated based on the probability distribution and the classification label. With the pre-training parameters of the visual Mamba model fixed, the cross-layer cue generator, the intra-layer cue generator, and the classification head are trained based on the loss function and the sample dataset until convergence, resulting in a trained optimized visual Mamba model. The selective prompting module processes input information based on the following steps: Obtain input information, which includes: an image token or the output of the corresponding Mamba block in the previous Mamba layer; The input information is processed using a cross-layer prompt generator to obtain cross-layer prompt information; The input information is processed using an in-layer prompt generator to obtain in-layer prompt information; Based on the fusion of cross-layer and intra-layer prompts using two dynamic scaling factors, the prompts are obtained. ; The prompt It is overlaid on the input information and used together with the category token as input to the corresponding Mamba block in the current Mamba layer.

2. The method as described in claim 1, characterized in that, The sample dataset contains diverse sample images and annotation information.

3. An image classification method, characterized in that, Image classification is performed using the trained visual Mamba optimized model as described in any one of claims 1-2.

4. An electronic device comprising a memory and a processor, the memory storing a computer program configured to be executed by the processor, the computer program including instructions for performing the method of any one of claims 1 to 3.

5. A storage medium storing a computer program, which, when executed by a computer, implements the method according to any one of claims 1 to 3.

Images