[Paper Review] DVPT: Dynamic Visual Prompt Tuning of Large Pre-trained Models for Medical Image Analysis

Along He, Kai Wang+3 more|arXiv (Cornell University)|Jul 19, 2023

Domain Adaptation and Few-Shot LearningComputer Science3 citations

TL;DR

This paper proposes DVPT, a parameter-efficient fine-tuning method for large pre-trained vision models in medical image analysis. By introducing learnable, dynamic visual prompts that adaptively attend to frozen, domain-transformed features, DVPT achieves state-of-the-art performance with only 0.5% of the model's parameters trainable, surpassing full fine-tuning by over 2.20% in Kappa score and reducing labeled data needs by up to 60%.

ABSTRACT

Limited labeled data makes it hard to train models from scratch in medical domain, and an important paradigm is pre-training and then fine-tuning. Large pre-trained models contain rich representations, which can be adapted to downstream medical tasks. However, existing methods either tune all the parameters or the task-specific layers of the pre-trained models, ignoring the input variations of medical images, and thus they are not efficient or effective. In this work, we aim to study parameter-efficient fine-tuning (PEFT) for medical image analysis, and propose a dynamic visual prompt tuning method, named DVPT. It can extract knowledge beneficial to downstream tasks from large models with a few trainable parameters. Firstly, the frozen features are transformed by an lightweight bottleneck layer to learn the domain-specific distribution of downstream medical tasks, and then a few learnable visual prompts are used as dynamic queries and then conduct cross-attention with the transformed features, attempting to acquire sample-specific knowledge that are suitable for each sample. Finally, the features are projected to original feature dimension and aggregated with the frozen features. This DVPT module can be shared between different Transformer layers, further reducing the trainable parameters. To validate DVPT, we conduct extensive experiments with different pre-trained models on medical classification and segmentation tasks. We find such PEFT method can not only efficiently adapt the pre-trained models to the medical domain, but also brings data efficiency with partial labeled data. For example, with 0.5\% extra trainable parameters, our method not only outperforms state-of-the-art PEFT methods, even surpasses the full fine-tuning by more than 2.20\% Kappa score on medical classification task. It can saves up to 60\% labeled data and 99\% storage cost of ViT-B/16.

Motivation & Objective

• Address the challenge of limited labeled medical data in fine-tuning large pre-trained models.
• Overcome the limitations of static prompt tuning methods that fail to adapt to complex texture and scale variations in medical images.
• Develop a parameter-efficient method that preserves domain-general knowledge while learning sample-specific, task-specific representations.
• Achieve high performance with minimal trainable parameters and reduced computational and storage costs in low-data regimes.

Proposed method

• Apply a lightweight bottleneck layer to transform frozen features to better capture domain-specific distributions in medical tasks.
• Introduce learnable visual prompts as dynamic queries that perform cross-attention with the transformed features to generate sample-specific representations.
• Project the attended features back to the original dimension and concatenate them with the frozen features to preserve pre-trained knowledge.
• Share the DVPT module across multiple Transformer layers to minimize the number of trainable parameters.
• Fine-tune only the DVPT module and task-specific layers while keeping the pre-trained backbone frozen.
• Control model capacity by adjusting the number of prompts and the number of shared layers.

Experimental results

Research questions

• RQ1 Can dynamic visual prompt tuning improve performance on medical image classification and segmentation with minimal trainable parameters?
• RQ2 How does DVPT compare to full fine-tuning and other parameter-efficient methods in terms of accuracy and data efficiency?
• RQ3 To what extent can DVPT reduce the need for labeled data in medical image analysis tasks?
• RQ4 Does the dynamic nature of the prompts enhance adaptation to complex texture and scale variations in medical images?
• RQ5 How does DVPT balance parameter efficiency, computation cost, and storage overhead in real-world deployment?

Key findings

• DVPT outperforms state-of-the-art parameter-efficient fine-tuning methods by more than 2.20% in Kappa score on medical classification tasks, even surpassing full fine-tuning with only 0.5% of the model's parameters trainable.
• With just 40% of the labeled data, DVPT achieves performance comparable to full fine-tuning with 100% data, demonstrating a 60% reduction in annotation requirements.
• DVPT reduces storage cost by up to 99% compared to full fine-tuning on ViT-B/16, while saving up to 10.22% in training memory.
• The method achieves the best trade-off between performance and model efficiency, with negligible increase in FLOPs and superior Dice and IoU scores on segmentation tasks.
• Optimal performance is achieved with 50 prompts, as higher numbers lead to optimization difficulties without performance gains.
• The DVPT module can be shared across multiple Transformer layers, significantly reducing the number of trainable parameters while maintaining high performance.