TransMed: Large Language Models Enhance Vision Transformer for Biomedical Image Classification

TL;DR

This study tackles adapting vision foundation models trained on natural images to few-shot clinical tasks with minimal annotated data. It introduces two key methods: a simple partial freezing fine-tuning strategy to balance learning capacity and data efficiency, and an LLM-based semantic guidance approach that contextualizes category labels to create discriminative semantic embeddings. On the MedFMC benchmark, partial freezing yields strong gains over standard fine-tuning and adapters, while LLM-contextualized labeling provides additional improvements, achieving 3–5% gains in 1-shot settings and clinching first place. The work demonstrates practical, data-efficient cross-domain adaptation of vision transformers for medical imaging when no medical pretraining data are available, with clear implications for rapid clinical deployment.

Abstract

Few-shot learning has been studied to adapt models to tasks with very few samples. It holds profound significance, particularly in clinical tasks, due to the high annotation cost of medical images. Several works have explored few-shot learning on medical images, yet they still require a large number of medical images for pre-training models to gain domain-specific priors. Vision foundation models recently have achieved remarkable success in natural images. Hence, adapting rapidly advancing vision foundation models from natural images to few-shot clinical tasks holds great promise. MedFMC has recently organized a challenge to shed more light on this topic at NeurIPS 2023. In this work, we present our challenge solution. We observe that a simple variant of fine-tuning with partial freezing shows remarkable performance. Empirical evidence demonstrates that this approach could outperform various common fine-tuning methods under limited sample sizes. Additionally, we explore enhanced utilization of semantic supervision to boost performance. We propose a novel approach that contextualizes labels via large language models (LLMs). Our findings reveal that the context generated by LLMs significantly enhances the discrimination of semantic embeddings for similar categories, resulting in a notable performance improvement of 3%-5% in 1-shot settings compared to commonly employed one-hot labels and other semantic supervision methods. Our solution secures the 1st place in the MedFMC challenge.

Figures (5)

• Figure 1: Performance overview on MedFMC in the 5-shot setting. Our approach significantly outperforms the standard adaptation of various foundation models, as well as the state-of-the-art few-shot learning technique, Visual Prompt Tuning vpt.
• Figure 2: Overall framework of our method. Given an image and a category name, we first utilize a large language model to generate contextual descriptions. Then we replace the category names in the description with [MASK] tokens and extract the embeddings of the [MASK] token using a masked language model as the semantic representation of the category name. We then align the visual embedding of the input image with the textual embedding by fine-tuning the vision foundation models with partial freezing.
• Figure 3: Emperical analysis of setting different numbers of frozen layers. We conduct evaluation on both ViT and Swin-Transformer under 1-shot and 5-shot settings across all datasets. Fine-tuning with partial freezing demonstrates a good compromise between full-model fine-tuning and linear probing, surpassing Visual Prompt Tuning (VPT).
• Figure 4: Performance changes on Colon 1-shot tasks as adaptation progresses. Fine-tuning with partial freezing demonstrated superior performance compared to other adaption techniques.
• Figure 5: The inter-class correlation matrices obtained by different language supervision methods on the ChestDR task. Left: Encoding category names. Middle: a template method using masked language models without contextualization. Right: our method leveraging large language models for label contextualization. It can be observed that the context generated by large language models plays a crucial role in fine-grained category distinguishing.