[Paper Review] - Merlin: A Vision Language Foundation Model for 3D Computed Tomography

lazyworm·2024년 7월 28일

PaperReview

목록 보기

1/1

그간 많은 Vision Language Model (VLM) 들이 제안되어 옴.

Medical 분야에서도 VLM 열풍이 불긴 했으나, 모두 2D 영상을 학습하는 것에 그쳤다.

근데 약 한 달전, CT를 위한 VLM이 arxiv에서 제안되었다...!!!

주의: 본 논문은 Main -> Results -> Discussion -> Method 순의 기이한 목차를 가지고 있음.

주의 2: 본 글은 논문의 실험결과를 상세히 다루지 않음. 상세한 실험결과는 논문을 참고할 것. (논문의 실험결과가 매우 상세하여 여기에 모두 기술하면 요약의 의미가 없어짐..)

Main

Problem Definition

Radiologist들이 판독해야 할 CT는 매해 6%씩 증가하여 판독 Burden이 증가하고 있음.
Medical VLM은 Radiologist들의 Burden을 줄여줄 수 있는 도구임.
그러나, 몇 가지 이유들로 인하여 Medical VLM들이 실전 투입되지 못하고 있음.
- EHR 코드, Radiology Report등 다양한 임상데이터들을 학습에 사용하지 않음
  - 특히, CT와 같은 3D Volume들이 VLM 학습에 사용되지 않음.
- 기존 Medical VLM들의 Downstream task의 임상 실효성이 낮음
  - 임상에서 실제로 필요하지 않은 task에 대한 성능 검증은 필요 없음!

Contribution

본 연구의 저자들은 Merlin이라는 CT를 위한 VLM을 제안함.
Merlin은 대규모 Pair data{CT, EHR 진단 코드, Report}로 학습됨
- 저 3가지가 한 쌍의 데이터임.
Merlin을 임상에서 쓰일만한 non-Adapted Task 3가지, Adapted Task 3가지를 설정하여 검증함.
Public에 데이터 및 모델 가중치 공개 예정. (24/07/28 기준 아직 공개 안됨.)

Discussion

저자들은 Merlin을 사용할 사람들을 위한 tip을 남겼다.
- Dataset을 가능한 많이 확보하여 학습 시킬 것.
- 높은 Image resolution의 CT를 사용할 것.
  - 더 정확히는, 학습시 Downsampling 하지 말 것. (선행 연구에서도 증명됨)
- 가능한 Batch를 키울 것.

필자의 생각도 아래에 추가로 정리해 보았다.
- 저자들이 논문의 제목에 Foundation이라고 언급한 것이 인상적임.
- 실제로, 후속 연구자들이 참고하기 쉽게 Merlin 실험결과를 매우 상세히 정리해 둠.
- 그러나, 역시 Foundation인 만큼 바로 임상에서 쓰기는 어려워 보임
  - 물론 높은 성능 수치이지만 아직까지는 성능 수치들이 아쉬운 수준임.
- 의사들은, Multi-modal로 환자들을 인식한다.
  - Merlin은 이제 시작이라고 생각한다. 이제 CT가 뚫렸으니 더 다양한 임상데이터들이 사용될 것으로 예상한다.
Inflated 3D를 사용한 것이 인상적임.
- 아무래도 Single GPU 및 3D input 지원을 위해서 선택한 것으로 보인다.
- I3D Backbone의 InfoNCE loss는 참신했음.
- 그러나 뭔가 좀 더 extensive한 pre-text task 실험이 필요해보임.

Method

Model

Architecture: Inflated 3D ResNet152 + clinical Longformer
- 2d weight compatibility를 위해, 2d weight를 3-channel로 stacking하여 사용
  - 단, 오직 conv layer인 경우에만 weight를 업로드 하도록 함.
- Longformer의 경우 context length를 최대 4096까지 지원함.
Pre-Training setup:
- Loss: BCE loss, InfoNCE loss
  - BCE loss: Phenotype classification
  - InfoNCE loss는 positive / negative pair를 구별하기 위한 loss.
- Optimizer: Adamw (lr = 1e-5, betas = (0.9.999))
- Etc: FP16 mixed precision, batch size 18 (GPU: A6000 48GB)
  - 단, staged 학습인 경우엔 Batch size를 22로 사용함.
  - 거의 대부분의 경우에서 staged 학습할 때보다 정식 Merlin 버전인 MTL이 더 성능이 좋으므로 무시해도 좋음..

Data (Pre-Processing, Split)

    - Preprocessing:
        - CT:
            - Resampling:
                - (x, y, z) 가 (1.5, 1.5, 3) spacing이 되도록 bilinear interpolation
            - Normalizing: -1000:1000 → 0:1
                - -1000 보다 작거나 1000보다 큰 intensity는 clipping
            - Patch size: (224, 224, 160) (Pad & Center crop)
        - EHR (Phenotype):
            - ICD9 / ICD10 코드 추출 후 PheWAS Phecode Mapping.
                - 총 16553개의 ICD 코드를 1692개의 진단 코드로 re-mapping
                - re-mapping 시에는 positive인 코드를 확인하여 tree label로 만듬
            - Mapping 된 코드들을 binarize하여 label로 변환
        - Radiology report
            - regex를 사용하여 report의 finding section만 추출함.
            - report가 길어질수록, 핵심 키워드에 overfit 될 것으로 추측함.
                - 이에 report를 anatomy 에 따라 split 하고, (위의 splitting.)
                    - train 시에 report의 각 부분을 달리해서 주거나 전체 report를 주는 hard-augmenting 사용 → 동일 anatomy에 대해 모델이 다양한 설명을 듣는 효과
                    - regex가 실패하는 경우엔 전체 report사용.
    - Split
        - Train: 15331 scans (60%)
        - Valid: 5060 scans (20%)
        - Test: 5137 scans (20%)

Results

1. non-Adapted Task

1-A. Zero-shot Findings Classification

Task Definition:
- 임의의 질환을 아무거나 골라서 Input CT 가 해당 질환에 대해 양성인지 음성인 판별.
Task Step:
1. 임의의 질환을 고름 (ex. Thrombosis)
2. 임의의 질환이 양성임을 시사하는 Prompt와 그 반대의 Prompt를 만듬
3. CT volume과 2에서 만든 Prompt들을 input으로 Merlin에 제공
4. Merlin이 해당 질환에 대한 양성 여부를 판별함.
이를 F1 Score로 성능평가함.

1-B. Zero-shot Phenotype (진단코드) Classification

Task Definition:
- CT input이 들어왔을 때 ICD Code grouping에 따른 phenotype을 분류해주는 task.
- 좀 더 정확히는 ICD-9, ICD-10의 코드가 사용되었음. (총 692가지 phenotype)
Results:
- 아래의 그림은 Phenotype의 prevalence와 AUROC를 Plot한 그림이다.

Task Definition:
- CT Image가 주어질 때, Matching 되는 Report 나 Impression Section을 찾는 task
- 혹은 Report 나 Impression Section을 받아 Matching 되는 CT Image를 찾는 task
- CT to Report / Report to CT에 대한 성능을 Recall@1로 측정함.
참고: Finding이나 Impression Section들의 단어 길이 통계.
- 평균적으로 Finding의 경우 21% 이상이 512개 이상의 token을 가진다.

2. Adapted Task

2-A. Multi-Disease 5-Year Prediction

Task Definition:
- CT Volume이 Input으로 주어질 때 해당 환자가 5년내로 만성 질환을 발병할 것인지 예측
- 만성 신장 질환, 골다공증, 심혈관 질환, 허혈성 심질환, 고혈압, 당뇨에 대해 수행함.
- AUROC로 성능을 평가함.

2-B. Radiology Report Generation

Task Definition:
- 말 그대로 input CT에 대한 report를 generation 하는 task.
- Encoder는 freeze하고, projection head와 RadLlama-7B를 fine tune하여 학습.
- 성능 측정은 아래의 Metric들을 사용함.
  - RadGraph-F1 [63], BERT Score [64], ROUGE-2 [65], and BLEU score[66]
Generation 결과는 human과 대체로는 일치하나, hallucination 현상도 존재함.