Fine-grained Prompt Tuning: A Parameter and Memory Efficient Transfer Learning Method for High-resolution Medical Image Classification

As illustrated in part (a) of Fig. 3, the FPT framework consists of two networks: a frozen LPM $M$ and a learnable side network $S$. Unlike other PETL methods that introduce additional learnable parameters within the LPM, in our approach, the entire LPM remains frozen while the side network is kept learnable and separate. We adopt a lightweight architecture design for the side network, which is a scaled-down variant of the LPM. The hidden dimensions of the side network are $1/k$ times that of the LPM, where $k$ represents a reduction factor. To leverage the pre-trained knowledge from the LPM, the side network reuses the intermediate features at each layer of the LPM. Specifically, given two models $M$ and $S$ with $L$ layers, parameterized by $\theta_{M}$ and $\theta_{S}$ respectively, the intermediate activation $z_{S,\mathrm{I}}^{l}$ of layer $l$ is obtained as follows:

where $\mathcal{F}$ denotes the module that fuses features of $M$ and $S$, and $z_{S,I}^{L}$ is considered the final output of the framework. The use of the side network not only reduces the number of trainable parameters due to its lightweight design but also helps mitigate memory expenses during the training phase. As shown in part (a) of Fig. 4, the side network eliminates the need for resource-intensive back-propagation from the pre-trained model by excluding any learnable parameters in the forward pass of the heavy LPM model.

Fine-grained information holds significant importance in the context of medical image analysis and is typically acquired by using high input resolutions. Although training with high-resolution inputs may be infeasible due to high memory consumption, using a LPM solely for inference with high-resolution images remains practical. Therefore, we propose an asymmetric input strategy within the FPT framework. Specifically, given a high-resolution image $I$, we simply resize it to obtain a low-resolution image $I^{\prime}$. Then, the frozen pre-trained model $M$ is provided with the image $I$, while the learnable side network $S$ is fed low-resolution image $I^{\prime}$. Thus, we have the intermediate activation of the side network $z_{S,I^{\prime}}^{l}=\theta_{S}^{l}(\mathcal{F}(z_{S,I^{\prime}}^{l-1},z_{M,I}^{l}))$.

In this section, we introduce our proposed fusion module $\mathcal{F}$. As shown in the part (b) of Fig. 3, we utilize the cross-attention mechanism [1] inside the FFM to fuse features from the LPM to the side network. In the context of cross-attention, we reuse the key and value from the self-attention layer of the pre-trained model $M$. Regarding the query, one approach is to directly reuse the query from the side network $S$. However, the cross-attention map can be large if the input sequence is long, leading to increased memory consumption. Therefore, we introduce a small set of learnable embeddings $z_{p}$, namely fine-grained prompts, into each layer of the fusion modules as the query. Unlike prompt tuning [13, 16], which directly uses prompts as part of the input sequence, fine-grained prompts serve as a bridge linking the frozen LPM and the side network. These prompts are concatenated with the intermediate sequence of the side network to join the forward propagation after fusing pre-trained features from the LPM. They are then removed after the layer’s forward processing. Specifically, the fusion module is processed as follows:

where the notations for the input source and the index of the layer are omitted in the formula for simplification. Here, $\text{CA}(\cdot)$ and $\text{Attn}(\cdot)$ denote the cross-attention module and attention function [25] respectively, $[\cdot,\cdot]$ denotes the concatenation operation, $f_{in}$ and $f_{out}$ denote linear layers that align the hidden dimension between the features. The terms $K_{\mathrm{M}}$ and $V_{\mathrm{M}}$ refer to the key and value mapping matrices within the corresponding self-attention layer of $M$ respectively.

Building on knowledge in the medical image domain, where images of the same modality often display similar anatomical structures and the objects of interest typically occupy a small proportion of the entire image, we propose a method to further reduce memory consumption, as shown in part (b) of Fig. 4. Specifically, we introduce important token selection, which selects the top $m\%$ tokens with the highest average scores on the self-attention map, considering them as important tokens. Only the features associated with these important tokens are passed to the FFM. This approach significantly reduces the overhead introduced by high-resolution inputs while preserving essential fine-grained information.

To further accelerate training procedure, we opt not to use any data augmentation on the input of the frozen LPM. This choice ensures that the intermediate features from the LPM associated with an image remain consistent throughout training. This approach allows us to pre-store these features before training, leading to significant reductions in training costs. Note that while the preloaded features remain fixed, the data augmentation applied to the input of the side network maintains the diversity of training samples.

We evaluate FPT on four medical datasets with different modalities, including fundus images (messidor-2 [4]), dermoscopic images (ISIC 2018 [3]), mammography (DDSM [15]), and chest X-ray (COVID [22]). The dataset sizes range from 1,748 to 11,527 samples, with classification categories varying from 3 to 7. We use official dataset splits when available. Otherwise, we employ a random partition of 70%/10%/20% for training, validation, and testing, respectively.

We compare FPT against full fine-tuning, linear probing, and state-of-the-art (SOTA) PETL approaches. In full fine-tuning, all parameters of the LPM are made learnable during training on the downstream tasks. Linear probing involves solely training the new task-specific head on top of the LPM. Generally, full fine-tuning often represents the upper performance bound for transfer learning, while linear probing represents the lower bound. We also compare FPT against five other popular SOTA PETL methods.

The performance and efficiency of PETL methods are tabulated in Table 1. It can be observed that fine-tuning the entire LPM (full fine-tuning), as the upper bound, achieves the best average AUC of 93.96%, requiring 24,116MB of training GPU memory with a batch size of 16. Although other PETL methods improve transfer learning efficiency by reducing the number of learnable parameters, they are unable to reduce the overhead brought by the high input resolution; all compared PETL methods require at least 20,000MB of memory, and even linear probing requires 4,364MB of memory. In contrast, FPT achieves the second-best AUC of 92.26%, while significantly reducing the memory requirement to only 3,182MB (13% of full fine-tuning). Moreover, with a separate network, FPT utilizes only 1.81% learnable parameters compared to full fine-tuning. Therefore, in terms of efficiency, FPT achieves the best PPE and PME, presenting the best trade-off between performance and efficiency. This parameter and memory efficiency positions FPT as a practical and feasible choice for leveraging LPMs in high-resolution contexts.

To assess the impact of the proposed components in FPT, we evaluate the performance and efficiency of our framework by incrementally incorporating components. As shown in Table 3, starting with a sole side network, although the side network is lightweight, the long input sequence still consumes a large amount of memory at 17,218MB for training. Employing fusion modules with fine-grained prompts to extract pre-trained knowledge from a frozen LPM notably enhances performance but further increases the memory burden. Then, the introduction of asymmetric input significantly lowers memory usage by 58% through decreasing the resolution of the input for the side network. It is worth noting that decreasing the input resolution for the side network enhances performance because the fine-grained information provided by the fine-grained prompts is sufficient, and smaller inputs for the side network reduce redundancy in features. Finally, by applying important token selection and preloading techniques, FPT further lowers the memory requirement by 64% without any loss of performance.

We evaluate the impact of different ratios of important tokens on FPT. As shown in Table 3, we observe that the classification performance remains similar when the ratio is between 20% and 50%. We then notice that 20% of the content of an image is sufficient for diagnosis in the tasks and modalities considered. Regarding efficiency, memory requirements increase as the ratio increases, leading to lower PME. Therefore, the ratio for important token selection is set to 20%.