VoCo: A Simple-yet-Effective Volume Contrastive Learning Framework for 3DMedical Image Analysis

논문 원본: https://arxiv.org/abs/2402.17300

VoCo: A Simple-yet-Effective Volume Contrastive Learning Framework for 3D Medical Image Analysis

 

논문 한 줄 요약: random sub volume이 어떤 공간적 prototype들과 얼마나 닮았는지를 예측하도록 학습 -> 3D 의료영상의 해부학적 위치 정보를 self-supervision으로 사용

---

1. Introduction 

기존 3D 의료영상 딥러닝은 높은 성능을 보이지만, 전문가 주석 비용이 매우 크다는 한계가 있다. 이를 해결하기 위해 자기지도학습(SSL)이 주목받아 왔으나, 기존 3D 의료영상 SSL 방법들은 주로 회전, 마스킹, 복원 기반의 프리텍스트 과제에 의존해 왔다. 이러한 접근은 저수준 시각 정보(texture, intensity)는 잘 학습하지만, 장기 단위의 고수준 정보를 충분히 학습하지 못한다는 문제가 있다. 실제로 선행 연구들은 사전 학습 단계에서 고수준 의미가 부족할 경우 다운스트림 성능이 크게 제한됨을 보여주었다.

저자들은 3D 의료영상의 중요한 특성으로 장기 간의 상대적 위치가 매우 일관적이라는 점에 주목한다. 예를 들어, 간·비장·신장·대동맥 등의 위치 관계는 환자 간에도 크게 변하지 않는다. 이러한 contextual position priors는 그 자체로 강력한 의미 신호이며, 이를 사전 학습에 활용하면 보다 의미적인 표현을 학습할 수 있다고 주장한다.

이를 바탕으로 본 논문은 contextual position prediction을 프리텍스트 과제로 사용하는 새로운 3D 의료영상 SSL 프레임워크 VoCo(Volume Contrast) 를 제안한다. VoCo는 공간적으로 분할된 볼륨을 prototype으로 삼고, 임의의 sub-volume이 어느 위치에 해당하는지를 볼륨 간 유사도 대비(volume contrast) 로 예측하도록 학습한다. 이를 통해 주석 없이도 고수준 의미 정보를 표현 공간에 주입하며, 다양한 다운스트림 작업에서 기존 SSL 방법들을 능가하는 성능을 달성한다.

 

2. Related Work

기존 자기지도학습은 주로 instance-level 대비 학습 또는 reconstruction 기반 학습에 의존해 왔다. Instance-level 대비 학습(SimCLR, MoCo)은 대규모 배치와 음성 샘플에 크게 의존하며, 이는 계산 비용이 큰 3D 의료 영상 환경에 적합하지 않다. 복원 기반 SSL(MAE 계열, rotate/mask/reconstruct)은 저수준 시각 정보 학습에는 효과적이지만, 장기 단위의 고수준 의미 정보를 충분히 학습하지 못한다.

 

Prototype-level 대비 학습은 이러한 한계를 완화할 수 있으나, 기존 방법들은 대규모 클러스터링이 필요하거나 무작위 초기화된 prototype에 의존해 의미적 안정성이 부족하다.

위치 기반 SSL(Jigsaw, Position prediction)은 공간 구조 학습을 시도했으나, 자연 영상에서는 객체 간 기하학적 관계가 일관되지 않아 효과가 제한적이며, 대부분 위치를 직접 회귀하는 방식에 머물러 있다.

 

VoCo는 이러한 한계를 극복하기 위해, 장기 간 상대적 위치가 일관적인 3D 의료 영상의 특성에 주목하여, 공간적으로 분할된 볼륨을 위치 기반 prototype으로 사용하고, sub-volume과의 volume contrast 를 통해 문맥적 위치를 예측한다. 이를 통해 별도의 클러스터링이나 좌표 회귀 없이도 고수준 의미 표현을 효과적으로 학습한다.

3. Method

VoCo는 두 개의 핵심 분기로 구성된다.

1. Prediction Branch: 임의로 크롭된 sub-volume이 전체 볼륨 내에서 어느 위치에 해당하는지를 예측
2. Regularization Branch: 서로 다른 위치를 대표하는 prototype(base)들이 의미적으로 충분히 구분되도록 정규화

이 두 분기는 동일한 backbone과 projector를 공유하며, 함께 학습된다.

 

Base crop과 Random crop

• 입력 3D 볼륨을 겹치지 않게 공간적으로 분할하여 nnn개의 base crop을 생성
  → 각 base는 특정 공간적 위치(문맥적 위치)를 대표
• 각 base crop을 backbone과 projector에 통과시켜 prototype feature qi 로 사용
• 동시에 볼륨 내에서 임의 위치의 random crop을 하나 생성하고, 이를 feature p로 변환

여기서 중요한 점은 random crop이 하나의 base에만 속하지 않고 여러 base와 겹칠 수 있다는 것이다.

 

Contextual Position Prediction (Prediction Branch)

• random crop feature ppp와 각 base feature qiq_iqi​ 사이의 cosine similarity 를 계산
• 이 similarity 값은 해당 sub-volume이 그 위치(base)에 속할 확률처럼 해석된다
• 실제 정답(label)은 random crop과 각 base 간의 공간적 겹침 비율로 생성된 soft label

이를 통해 VoCo는 좌표를 직접 회귀하거나 분류하지 않고, 어느 위치와 얼마나 닮았는가를 예측하도록 학습한다.

이 예측을 감독하는 손실이 Lpred 이다.

 

Volume Contrast Regularization (Regularization Branch)

단순히 위치 예측만 하면, base들이 서로 비슷해지는 collapse 문제가 발생할 수 있다.
이를 방지하기 위해 VoCo는 추가적인 정규화 손실을 도입한다.

• 모든 base prototype qiq_iqi​ 쌍에 대해 cosine similarity를 계산
• 서로 다른 base 간 유사도를 0에 가깝게 유도
• 결과적으로 base들은 서로 직교에 가까운 방향을 갖게 됨

이 정규화 항이 Lreg -> 각 base가 서로 다른 의미적 위치(장기 문맥) 를 명확히 대표하도록 만든다.