[22’ ECCV] BioViL: Making the Most of Text Semantics to Improve Biomedical Vision-Language Processing
카테고리: Medical
🔍 Abstract
이 논문은 지금까지의 biomedical text embedding이 충분한 domain-specific semantic representation을 학습하지 못했음을 지적하며, (1) 새로운 text embedding model로 CXR-specific LM인 CXR-BERT를 제안한다. 더 나아가 이를 활용하여 CLIP과 비슷한 (2) VLP(Vision-Language Pretraining) 모델인 BioViL을 제안한다. 또한 (3) CXR Phrase Grounding Benchmark로 MS-CXR을 구축하였다. 이 모델은 CXR domain의 Biomedical Vision-Language task에서 state-of-the-art 성능을 보여준다.
1. CXR-BERT
CXR-BERT에서는 기존 BERT 학습에 사용되던 MLM(Masked Language Modelling) Loss에 RSM(Radiology Section Matching) Loss라는 새로운 objective를 추가한다. 이는 저자들이 radiology report를 보면서 떠올린 새로운 contribution이다.
Radiology report는 일반적으로 Findings과 Impression의 section으로 나누어져 있고, Findings에서는 CXR을 보고 어떤 부분에 문제가 있는지를 기술하며, Impression에서는 이를 요약하여 진단을 내린다. 따라서 이 두 section은 서로 밀접한 관계를 갖고 있으며, 이를 이용하여 RSM Loss를 정의한다.
저자들은 기존에 BERT가 text classification을 위해 [CLS] token을 사용하는 것을 참고하여, Findings와 Impression의 section을 각각의 [CLS] token으로 표현하고, 이 두 token 간의 cosine similarity를 최대화하는 방향으로 학습한다. 즉, Findings 부분의 word tokens $\mathbf{w}^ {F} = (w_ 1, \cdots, w_ T)$와 Impression 부분의 word tokens $\mathbf{w}^ {I} = (w_ 1, \cdots, w_ T)$에 의해 BERT로부터 [CLS] token embedding $\tilde{\mathbf{t}}^ {F}$와 $\tilde{\mathbf{t}}^ {I}$를 얻는다. 그리고 Text Projection Layer, 즉 일종의 MLP인 $P_ {\text{txt}}$를 통과하여 얻은 두 vector $\mathbf{t}^ {F}$와 $\mathbf{t}^ {I}$를 사용하여 InfoNCE Loss를 계산한다.
Total loss는 이를 weighted sum으로 처리한 $\mathcal{L} = \mathcal{L} _ {\text{RSM}} + \lambda_ {\text{MLM}} \mathcal{L} _ {\text{MLM}}$이 된다. 이처럼 Radiology report에서만 사용할 수 있는 새로운 objective를 추가하여 CXR-specific LM을 학습할 수 있었다.
2. BioViL
BioViL(Biomedical Vision-Language Model)의 전체 구조는 위 그림과 같다. 이는 CLIP의 구조와 크게 다르지 않은데, 다른 점이라면 1) Language Model의 성능을 유지하기 위해 MLM(Masked Language Modelling) Loss를 추가로 사용하고, 2) Local Feature에 대한 정보를 확인하기 위해 명시적으로 Visual Feature를 Spatial Location을 보존한 채로 encoding한다는 것이다. 각각의 과정은 다음과 같다.
- Image Domain: ResNet-50 Encoder $E_ {\text{img}}$를 사용하여 local image embedding $\tilde{\mathbf{V}}$를 얻는다. 이후 Two-layer projection model $P_ {\text{img}}$를 통과하여 $\mathbf{V}$를 얻는다. 이는 차원을 줄이는 역할을 한다. 마지막으로 Mean pooling을 통해 global image embedding $\mathbf{v}$를 얻는다.
- Text Domain: CXR-BERT $E_ {\text{txt}}$에서의 [CLS] token을 사용해 text embedding $\tilde{\mathbf{t}}$를 얻는다. 이후 Two-layer projection model $P_ {\text{txt}}$를 통과하여 $\mathbf{t}$를 얻는다.
학습 시에는 Local Alignment 등은 계산할 만한 데이터셋이 없으므로 Global Alignment만을 사용한다. 이는 CLIP의 objective인 InfoNCE Loss와 동일하다.
실제 Total loss는 $\mathcal{L} = \lambda_ {\text{GA}} \mathcal{L} _ {\text{GA}} + \mathcal{L} _ {\text{MLM}}$로 정의된다. 이는 Global Alignment Loss와 MLM Loss의 weighted sum이다. 학습 시에는 MIMIC-CXR v2(147K) 데이터셋을 사용하였다.
3. MS-CXR
저자들은 CXR Phrase Grounding Benchmark로 MS-CXR을 구축하였다. 이는 CXR 이미지의 특정 부분과 radiology report를 매칭하는 task로, 이를 통해 성능을 확인할 뿐만 아니라 모델의 interpretability를 확인할 수 있다.
4. Experiments
BioViL의 Evaluation은 (1) Classification (Zero-shot, Fine-tuning), (2) Segmentation에 더하여 새로 추가된 task인 (3) Phrase Grounding을 확인하였다.
4.1. Classification & Segmentation
두 부분에서 전부 좋은 성능을 보여주었다. 이는 결국 CXR-specific LM인 CXR-BERT의 강력함을 보여주는 것이다. 즉, text input을 잘 분석하여 이들의 특성으로부터 학습을 진행한 것이 성능 향상에 큰 도움이 되었다고 볼 수 있다.
4.2. Phrase Grounding
Phrase Grounding은 text와 유사한 local image feature를 찾는 것이므로, 위와 같이 높은 유사도를 가지는 부분이 diagnosis에 핵심적인 부분임을 파악할 수 있다. 이는 interpretability 측면에서도 큰 도움이 된다.
결과적으로 BioViL은 이러한 task에서 다른 모델들을 큰 폭으로 앞서는 성능을 보여주었다. 단, GLoRIA의 경우 local feature attention을 이용하여 성능을 향상시킨 모델이었기 때문에 기존 모델 중에서는 가장 높은 성능을 보여주었다. 따라서 저자들은 BioViL에 GloRIA의 local loss term을 추가하여 학습하였고, 이를 BioViL-L이라 명명하였다. 이러한 방식은 아주 효과적이어서, Phrase Grounding에서 더 높은 성능을 보여주었다. 아래는 이러한 GLoRIA의 local loss 계산 방식을 나타낸 그림이다.
4.3. Ablation Study
참고로, Ablation Study로부터 한 가지 좋은 결론을 얻을 수 있었다. CLIP과 같은 Contrastive Learning을 수행하는 동안 Language Model을 MLM(Masked Language Modelling) Loss로 추가로 학습하는 것이 성능 향상에 큰 도움이 된다는 것이다.
💡 Summary
CXR-BERT는 Radiology Report가 Findings, Impression Section으로 나누어져 있으며 이 둘은 밀접한 관계가 있다는 사실로부터 RSM(Radiology Section Matching) Loss를 정의하여 학습하였다. 이를 통해 CXR-specific LM을 얻을 수 있었다.
BioViL은 이러한 CXR-specific LM을 활용하여 구축한 Biomedical Vision-Language Model으로, Ablation Study를 통해 이러한 domain-specific constrastive learning 시에는 language model을 추가로 학습하는 것이 성능 향상에 큰 도움이 된다는 것을 확인하였다.
마지막으로 저자들은 MS-CXR이라는 CXR Phrase Grounding Benchmark를 구축하여 모델의 성능을 확인하였으며, 이때 local loss term을 추가하여 학습하는 것이 phrase grounding에 큰 도움이 된다는 것을 확인하였다.
저자들이 제안한 BERT 학습 방법은 CXR이 아닌 다른 domain에서도 쉽게 적용할 수 있는 방법은 아니다. 대신, 이처럼 domain-specific data를 잘 분석하여 모델에 녹여내는 과정은 어느 domain에서나 중요하다는 점을 기억해두자.
댓글 남기기