들어가며..

요즘 diffusion 계열의 generative model 보다는 CLIP 계열의 representation alignment 방식이 self supervised learning 에 더 유리하다는 생각을 한다. 다만, CLIP 은 입력 영상을 global pooling 을 하여 하나의 representation vector 으로 만들어 버리는데, chest CT 처럼 영상들이 다들 비슷하게 생겼을 때 이 방법은 구석 구석 숨어있는 디테일을 모두 압축시켜버리기 때문에 적합하지 않다.

영상을 하나의 벡터로 압축해버리지 않고 각각 영역의 embedding 을 보존할 수 있는 방법이 없을까?

고민을 하고 논문을 찾아보았는데, 너무 당연한지, 이 생각을 모두 하고있었다.

fVLM 논문에서 발췌한 그림인데, CLIP 은 heatmap 을 만들 때 영상 전체를 참고하여 localization 이 전혀 안된다는 내용이다. 그럼 어떻게 할까? DAMO Academy 에서는 organ segmentation 을 활용하여 organ 별로 contrastive learning 을 해야 한다고 주장한다. 복잡한 Chest CT 영상을 divide-and-conquer 하겠다고 생각해도 좋다.

이 그룹에서 CT-GLIP, fVLM 논문 (and more..) 을 찍어냈는데, 이 아이디어로 여러 편 우려먹을 생각인가보다..

CT-GLIP

Motivation

• CLIP 은 image 와 text 의 global alignment 에 의존함
• Chest CT 영상은 information 이 굉장히 sparse 하고 복잡해서 one representation 으로 collapse 시키면 정보가 너무 많이 사라짐
• GLIP 은 anatomical region 과 textual findings 를 pair 할 수 있음

Methods

Architecture

• Dual-path architecture: CLIP 이랑 똑같고, vision encoder 에는 CNN (nnUNet) / ViT (MiT), text encoder 에는 ClinicalBERT 를 활용
• Organ-level average pooling: TotalSegmentator 을 활용해서 얻은 organ 의 segmentation 을 "feature map with highest resolution" 에 맞게 downsampling 하여 average pooling 을 수행함. 근데 이게 encoder 쪽인지 decoder 쪽인지는 모르겠음.
• Additional two-layer MLP: Pooling 이후에 projection
• Frozen text encoder: catastrophic forgetting 방지

Results

Anatomy contrastive learning

Purpose: grounded visual embeddings 가 organ category 와 매칭이 잘 되었는지 확인하는 작업.
Method: Image 는 3D encoder 에 넣고 organ mask 를 활용해서 organ-level pooling 을 수행하고, Text 는 "this is the {organ_name} in the CT scans" 로 작성하여 embedding 을 구한다. 각 organ 별로 similarity 가 제일 높은 text pair 을 찾아서 classification 을 수행한다.

• Segmentation loss 도 추가했음

그 결과로 CLIP 은 0.01%, GLIP 은 86.9% 를 달성했다고 하는데... 아니 organ-level pooling 을 넣어줬으니 가능한거 아닌가? 오히려 90%도 안된다는게 놀라움.

Zero-shot Abnormality Detection

Purpose: 어떤 organ 에 병이 있는지 binary classification
Method: 위랑 비슷한데 text 에 a pair of normal and abnormal text descriptions 를 넣어서 어떤 prompt 와 similarity 가 더 높은지 비교함.

• 언급 없는 organ 은 "no evident abnormality in {organ}"
• negative pairs 를 늘리기 위해 abnormality dictionary 를 만드는데, 104개의 organs 에서 diverse 한 abnormality descriptions 를 저장해서 훈련중에 불러옴.

Fine-tuning

• Per-organ tumor detection / segmentation 을 수행했다.

Scratch < CLIP < GLIP 순으로 성능이 좋았다. Breast, esophagus 에서는 GLIP 의 성능이 조금 더 떨어졌는데, tradeoff 가 있나?

Comments

• lung 은 그냥 하나로 퉁쳐버리고 작업했다. Lung 도 엄청 큰 기관인데 이 안에서는 따로 분할된 데이터를 정제하기가 까다로웠다나.
• 그래서 결국 organ 안에서의 localization 이 불가능하다.
• 다만, 어떤 organ 이 문제가 있는지 알려준다는 것은 결과적으로 긍정적인 pipeline 이라고 생각한다.

fVLM

CT-GLIP 에서 뭐가 맘에 안들었을까?

1) Global average pooling 을 사용한 것
2) 단순히 patient level 로 학습하면 false negative 가 너무 많음

GAP -> Self attention

사실 이 부분은 논문에서 딱히 ablation study 로 발전시키지는 않았는데, CT-GLIP 에서는 global average pool 해놓고 여기서는 token 을 추출한 뒤 self attention 을 수행한다. 이게 맞는 것 같은게, 병은 전반적으로 퍼진 diffuse 한 병과 종양같이 focal 한 병이 존재할 수 있다. 단순히 global average pooling 을 했다가는 이런 신호가 dilute 될 수 있기 때문에 좋은 발전이라고 생각한다.

이 피규어가 처음에는 사실 잘 눈에 안들어왔었는데, 세 번 정도 읽어보니 알 것 같다.

1. ViT 를 활용해서 patchify 한다
2. 각 organ 의 segmentation foreground mask 를 포함한 patch 들을 flatten 한다 (organ 의 크기가 다르니 각각 patch 의 수가 다를 것이다!!)
3. 맨 뒤에 organ-specific query token 을 하나 붙여준다
4. Anatomy-wise self attention 을 통해 query token 을 update 해준다
5. Update 된 query token 을 활용해 contrastive learning 을 수행한다!

False-negative problem

Chest CT 에서는 positive 라고 간주해도 되는데 단순히 같은 pair 에서 나오지 않았다고 해서 negative 로 취급되는 경우 두 가지가 있다.

1) 정상일때

영상 A 의 abdomen 이 normal 이고, 영상 B 의 abdomen 이 normal 이라면, 이 둘이 같은 영상에서 나오지 않았다는 이유만으로 similarity 를 0으로 만들기에는 꽤나 비슷하게 생겼을 것이다 (사람의 장기는 비슷한 형태를 띄고 있기 때문이다). 물론 그 장기 안에서 또 사람마다 너무 차이가 난다면 다시 생각해봐야겠지만, 일단 chest CT 에서는 정말 많은 장기들이 존재하기 때문에 충분히 고려할만한 상황이다.

해결: 어떤 서로 다른 영상의 장기들이 서로 정상인 경우 label 을 1 으로 수정한다.

2) 둘다 정상은 아니지만, 똑같은 병이 있을 때

물론 사람마다 병이 다르게 생기긴 했지만, 전형적인 병의 형태가 있지 않은가. 이 논문에서는 같은 병이 있을 때는 similarity 가 조금 높아도 된다고 주장한다. 다만, 둘다 정상일 때 처럼 무지성으로 label 을 1로 바꿀 수는 없다. 병의 위치와 정도가 편차가 크기 때문이다.

해결: Bootstrapping 을 활용한다. Co-teaching scheme 을 활용하는데, fVLM 모델 두개를 동시에 훈련하며 나온 prediction 을 서로의 label 로 update 해준다.

최종 label

1과 2 평균때린다 ㅎㅎ

Results

CLIP 이랑 GLIP 이긴다.. 말고는 딱히 뭐 없다.

Ablation study 로 pooling 을 바꿔줬으면 좋았을 것 같은데, 이게 좀 아쉽다.

Conclusion

fVLM 을 마무리로 chest CT 를 organ level 로 나누어서 divide-and-conquer 하는 접근은 연구가 끝났다고 생각한다.

다만, organ 안에서의 complexity 를 충분히 modeling 했는지는 여전히 의문이 든다. 더 좋은 localization 방법론은 없을까?