
Guiding Medical Vision-Language Models with Explicit Visual Prompts: Framework Design and Comprehensive Exploration of Prompt Variations (NAACL 2025, Zhu et al.)
본 논문은 medical-vqa task를 해결하기 위해 explicit한 visual prompt를 제안하는 논문으로, 몇 주 전에 accept 되어 올해 NAACL에 publish 될 따끈따끈한 논문이다.
Medical-VQA 태스크를 해결하기 위해, 많은 모델들이 Vision-Language Model, VLM을 활용한다. BLIP, CLIP, LLaVA와 같은 모델 아키텍쳐를 따르며, medical domain만이 보유하고 있는 이미지와 질문의 특성을 학습하기 위해 medical domain alignment를 진행한다. 보통 (image,text) 쌍을 지어 large-scale dataset을 구축하고, pre-trained VLM을 finetuning한다. 그렇게 하여 학습된 Medical-domain aligned VLM 모델을 기반으로 SLAKE, PathVQA, VQA-RAD와 같은 VQA 데이터셋에 파인튜닝하여 이미지와 질문이 주어졌을 때 답을 생성하도록 한다.
VLM 모델들은 이미지와 텍스트 데이터를 처리하는 데 효율적이지만, 한 이미지에서 중요한 정보를 추출하는 데에는 한계가 존재한다. Pretraining 단계에서 image content에 대한 설명을 생성하는 과정을 주로 채택하기에, 특정 specific region보다는 global한 visual feature에 집중하게 되는 경향이 있다. 이 과정에서 특정 region이 guide로 주어지지 않고, attention mechanism을 활용하여 implicit한 attention 분포를 자동적으로 학습하게 된다. 그러나 VQA나 Referring Expression Comprehension (REC)와 같은 테스크에서는 특정 물체나 localized된 지역을 보는 것을 요구한다.
이를 위해, visual prompt를 explicit 혹은 implicit하게 활용하는 연구들이 진행되고 있는 추세이다. 해당 논문은 Medical VQA 태스크를 수행하기 위해, explicit한 visual prompt를 활용하는 방법론, MedVP를 제안한다. MedVP는 ViP-LLaVA를 Medical VQA에 적용하는 법에 대해 연구한 것으로 볼 수 있다.
시각적 프롬프트를 활용한 ViP-LLaVA에 대한 간단한 설명은 다음과 같다. ViP-LLaVA는 CVPR 2024년 논문으로, 이미지 자체에 marker등으로 집중할 영역을 표시하고, Text Prompt로 "What is the person marked with the red arrow holding"과 같이, 집중 영역에 대한 답을 생성할 수 있도록 한다.
### 2. MedVP Method
MedVP는 크게 총 3단계로 이루어진다. 1단계에서는 ROI를 나타내는 이미지와 Medical Report를 기반으로 VLM을 학습시킴으로써, medical domain alignment를 진행한다. 그리고 2단계에서는, VQA dataset에 finetuning하기 이전에, 질문과 관련있는 ROI 영역을 추출하는, visual grounding 작업을 진행한다. 예를 들어, 질문에서 heart에 대해 물어보면 이미지에서 heart 부분을 추출하여 이미지에 표시하는 것이다. 마지막으로, 2단계에서 추출한 local 영역과 그에 대응하는 visual prompt의 표시 (red ractangle, blue ellipse)를 기반으로 QA Task에 finetuning하는 작업이다.
MedVP의 첫 단계는 LLaVa-MED에서와 같이 medical 도메인에서의 content와 vision knowledge를 alignment를 진행하는 과정이다. 본 논문에서는 medical iamge에서는 진단을 진행할 때 구체적인 디테일이나 부분적인 지역 (legions or organs)을 바라보는 것이 중요하다고 강조하며, global 정보와 같이 key region을 이해할 수 있도록 학습을 진행한다. 이를 위해, local region에 대한 annotation을 포함한다. 일반적으로 (image, text) 두 쌍의 pair만을 활용한다면, MedVP는 (I: image, P: 이미지 내 ROI를 나타내는 visual prompt, T: text) 세 쌍을 묶어 활용한다. I와 P를 결합하기 위해서는 ViP-LLaVA의 방식을 따르는데, 이는 I와 P를 alpha blending을 활용하여 전체 이미지와 관련있는 local 지역을 묶는 것으로 다음과 같이 표현할 수 있다.
는 0과 1사이의 숫자로 visual prompt의 transparency level을 나타낸다.
최종적으로, local and global 특성을 지니는 blended image를 기반으로 autoregressive language modeling을 활용하여 ground truth text 를 맞추기 위해 생성 토큰의 likelihood를 maximize한다.
학습을 위해서는 (image, ROI, text description)을 내포하는 MedTrinity-20M 데이터셋의 일부를 활용한다.
다음 단계는 어떤 question이 주어졌을 때, 그에 대응되는 local region을 찾기 위해 수행되는 과정이다. 한 이미지에 대해 많은 질문들이 나올 수 있고, 그 질문마다 중요한 local region이 다를 것이기에, 해당 과정이 manual annotation 없이는 어려운데, 본 논문에서 manual하게 진행하지만 explicit한 visual prompt를 활용하기 위한 프레임워크를 잘 짰다고 생각한다. 우선 질문에 대해 중요한 region-level entity를 추출하고 그 다음에 해당 entity를 visual prompt extractor를 활용하여 visual prompt를 생성하게 된다.
우선 entity를 추출하기 위해 LLM 모델을 활용한다. 각 질문 Q에 대해 LLM이 entity set을 추출하도록 한다.
본 연구에서 활용된 prompt는 아래와 같다.
You are a medical entity identifier. Please help me find the medical entity that could help visual question answering if the entity is cropped by bounding box in the corresponding image. Basically you need to identify the organ and disease.If any specific organ is mentioned in the question, include the unmodified noun in your response. Otherwise, add exact 'organ' in your response if that labeling can be helpful.
If you think labeling a disease would help answering the question, add exact 'disease' in your response. Here are some example sentences and answers: Sentence: Does the kidneys look abnormal? Answer:
['kidney', 'disease']. Sentence: Is there swelling of the grey matter? Answer: ['disease'].
Sentence: What is the largest organ in the picture? Answer: ['organ"]. Sentence: Which is the biggest in this image, lung, liver or heart?
Answer: ['lung', liver', heart']. Sentence: Which organs appear in pairs? Answer: ['organ'].
Following the example, answer the question:
Sentence: [SENTENCE]
윗 단계에서 question에 해당하는 region-level entity를 추출한 후, visual prompt extractor를 활용하여 이미지에서 해당 위치 (visual prompt coordinates)를 찾아내는 작업이다. 이를 위해서는 Grounding DINO, (image, text) input을 활용한 detection 모델을 활용한다. Grounding DINO는 general domain image에서 open-vocab detection에서 좋은 성능을 보인 모델로, 모델 아키텍쳐를 아래에 나타낸다.
Grounding DINO가 좋은 성능을 보일지라도, general domain에 pretrain된 모델이라, medical image domain에서 특정 entity에 대한 region을 확실히 구하는 것이 어렵기에, 본 논문에서는 해당 모델을 medical image domain에 finetuning하여 활용한다. 파인튜닝 과정에서는 classification을 위해서 object와 language token 사이의 contrastive loss와 bounding box regression을 위한 L1과 GIoU loss를 활용한다. 파인튜닝 과정이 마치면, 이미지와 2.2에서 도출한 entity를 Grounding DINO에 넣어 관련있는 region의 visual prompt를 도출할 수 있게 된다.
$$
\text{P}= \text{G-DINO}(\text{I}, \mathcal{E})
$$
학습을 위해서는 SLAKE와 SA-Med2D dataset의 일부를 활용하였다.
마지막으로 위에서 도출한 visual prompt를 활용하여 answer를 generate하는 과정이다. 다양한 visual prompt에도 답을 생성하는 모델을 구축하기 위해 scribble, rectangle, ellipse와 같은 다양한 시각적 요소를 활용하여 2.1에서 언급한 alpha blending을 수행한다. 그리고 모델이 주어진 시각적 프롬프트에 집중할 수 있도록 color와 category를 visual annotation으로 활용하여 MedVP에 넣어준다. 예를 들어 빨간 직사각형과 파란 타원이 있다면 다음과 같이 LLM (2.1- medical align된 MedVP)에 넣어준다.
Question: Where is the liver in the image?
Inserted Prompt: Please pay attention to the areas within the red rectangle and the blue ellipse, they may contain useful information for generating the question.
Question: Where is the liver in the image?
Downstream Finetuning을 통해 Medical VQA에서의 성능을 확인하기 위해 총 3가지 데이터셋, SLAKE, VQA-RAD, PMC-VQA를 활용하였다.
Base model로 ViP-LLaVA 7B를 사용하였고, RTX 4090과 H100 GPU를 통해 학습을 진행했다고 한다. 2.1 Region level medical domain alignment를 위해서는 lr: 2e-5 batch size: 64, Groudning DINO finetuning 시에는 prediction score threshold: 0.2, batch size : 24, 그리고 downstream fine-tuning 단계에서는 lr: 2e-5, batch size 128, warmup ratio 0.03을 활용하였다고 한다.
Downstream Finetuning 과정은 4개의 RTX 4090으로 SLAKE 6 hours, VQA-RAD 2 hours, PMC-VQA 20 hours가 걸렸다고 한다. Groudning DINO finetuning 과정도 동일하게 4개의 RTX 4090으로 약 9시간 걸렸다고 한다. Medical domain alignment를 위해 수행된 시간을 적혀있지 않은 것으로 보이지만, H100을 해당 과정에서 활용하지 않았을까 하는 유추를 한다.
아래 테이블은 downstream VQA task에서의 성능으로 이전 논문들에 비해 월등히 향상된 성능을 보인 것으로 보고 되었다.
아래 테이블은 Visual Prompt가 있고 없고에 대한 ablation study로, visual prompt를 통해 Medical VQA의 context 내에서 local visual information을 capture하는 것에 도움이 됨을 나타낸다.
아래 사진은 cross-attention map의 visualization 결과로, 노랑색 부분이 모델의 attention value를 나타내는데, visual prompt의 영역과 많이 overlap하는 것을 보며, visual prompt를 통해 모델이 해당 부분이 집중하여 학습을 진행할 수 있었음을 알 수 있다.
본 논문에서는 모델이 natural하고 robust하게 relevant region에 대응하는 학습을 수행하였는지 확인하기 위해, limited visual prompt가 있을 때의 결과도 보고한다. 각 테스트 데이터셋 별로, Groudning DINO가 생성하는 visual prompt coordinate의 일부 (80~20%)만을 활용하였을 때 각각의 결과이다. Visual prompt를 줄일수록 result가 살짝 감소하는 것으로 보이지만, 전체적인 성능이 stable한 것으로 나타났다. 저자들은 이는 visual prompt가 없는 경우에도 explicit attention guidance를 통해 꽤 유사한 ROI에 대응한다고 보며, visual prompt가 직접적으로 model의 attention을 주입시키는 것이 아닌, relevant region에 focus하도록 유도한다고 해석한다.
본 논문에서 visual prompt가 실제로 답변을 잘못 생성하는 경우도 있는데, 몇 가지 예시를 남긴다
(a) Visual Prompt가 너무 작은 경우 모델에서 인식하기에 어려워 잠재적으로 error를 발생시킬 수 있음
(b) Counting question에서 실제 질문과 관련없이 visual prompt 수에 의존하여 잘못된 정답을 내뱉는 경우
(c) Existence of abnormalities과 관련된 질문에서 abnormality의 존재 여부와 관련 없이 visual prompt의 존재 여부로 답변하는 경우
(d) Location과 관련된 질문에서 visual prompt의 위치 자체를 답변으로 내는 경우가 존재할 수 있다고 한다.
해당 논문이 medical vqa에서 question에 대응되는 image이 중요한 lesion과 organ 같은 위치를 찾는 것의 중요성을 알리는 데에 좋은 기여를 했다고 생각한다. 실제로 VLM에 맡겨버리는 모델들이 많은 반면에, 본 논문은 실제로 QA가 진행되는 방식을 고려하였다는 점에서 의료 분야에서 널리 활용될 수 있는 방식인 점에 높게 평가한다. 다만, inference 과정에서 entity 추출을 위한 LLM 모델 하나, region recognition을 위한 Groudning DINO 모델, 그리고 최종적으로 답변 생성하는데 VLM까지 도입하면서 필요한 모델의 양이 많아 computation cost 측면에서는 부담이 되는 부분도 있다고 생각하지만, 우선 explicit하게라도 실제 의사가 답을 내뱉는 과정을 resemble 했다는 측면에서 좋은 초기 연구로 자리 잡을 것 같다.