기술 요약: 3D CT-GraphEdgeMAE

문제 정의

컴퓨터 단층촬영(CT) 이상 탐지는 악성 종양의 조기 진단에 매우 중요하지만, 현재의 접근 방식들은 상당한 난관에 봉착해 있다. 지도 학습 기반의 딥러닝 방법(예: 3D U-Net 변형 모델)은 노동 집약적인 복셀 수준의 주석(annotation)에 크게 의존하며, 이는 임상 현장에서 구하기 매우 어렵다. Masked Autoencoders (MAE) 및 Vision Transformers (ViT)와 같은 자기지도 학습 방법은 주석의 필요성을 줄여주지만, 3D CT 데이터에 적용할 때 높은 계산 비용과 메모리 요구량을 초래한다. 구체적으로, ViT와 MAE의 어텐션 메커니즘은 패치 수에 따라 이차적($O(N^2)$)으로 증가하여 고해상도 3D 볼륨 적용 시 "계산 폭발" 현상을 일으킨다. 또한, 이러한 패치 기반 접근 방식은 3D 결절에 내재된 엣지 침윤(edge infiltration)이나 종양과 주변 조직 간의 공간적 연속성과 같은 명시적인 국소 관계 및 위상 구조를 포착하는 데 실패하는 경우가 많다. 저자는 기존의 그래프 기반 방법(예: GraphMAE)이 노드(node)를 마스킹할 뿐 에지(edge)를 마스킹하지는 않으며, 이로 인해 결절 탐지에 필수적인 복셀 블록 간의 관계를 명시적으로 모델링하지 못한다고 주장한다.

방법론

본 논문은 3D CT-GraphEdgeMAE를 제안한다. 이는 소규모 데이터셋 및 자원이 제한된 환경을 위해 설계된 엔드 투 엔드(end-to-end) 자기지도 학습 프레임워크이다. 이 방법은 3D CT 볼륨을 그래프로 변환하며, 여기서 복셀은 노드가 되고 공간적/특징적 유사도가 에지를 정의한다. 프레임워크는 다섯 가지 핵심 구성 요소로 이루어져 있다:

1. 3D CT 그래프 구축:

   • 노드: 2D 패치나 사전 분할된 해부학적 구조를 사용하는 대신, 이 방법은 3D 복셀을 노드로 취급한다.
   • 특징: 노드 특징은 통계적, 위치적, 질감, 위상적, 학습 가능 및 확장 가능한 특징을 결합한 하이브리드 전략을 통해 초기화된다.
   • 에지: 에지는 공간적 거리와 특징적 유사성을 기반으로 구축되어 CT 스캔의 공간적 볼륨 구조를 보존한다.

2. 에지 마스킹 전략:

   • 노드를 마스킹하는 GraphMAE와 달리, 3D CT-GraphEdgeMAE는 에지를 마스킹한다. 이는 결절이 공간적 그래디언트 특징을 가진 여러 복셀에 걸쳐 존재한다는 관찰에 근거한다. 에지를 마스킹함으로써 모델이 단순히 블록 자체를 배우는 것이 아니라, 복셀 블록 간의 관계를 명시적으로 학습하도록 강제한다.

3. 인코더 (GCN):

   • 인코더로 그래프 합성곱 네트워크(GCN)를 사용한다. 저자는 동적인 인접성 계산과 관련된 계산 오버헤드 및 과적합 위험을 피하기 위해 그래프 어텐션 네트워크(GAT) 대신 GCN을 선택했다. GCN은 여러 홉(hop)에 걸쳐 점진적이고 암시적인 특징을 학습하여 에지 예측을 위한 문맥적 관계를 제공한다.

4. 디코더:

   • 디코더는 두 노드($v_i, v_j$) 사이의 마스킹된 에지의 존재 여부를 예측한다.
   • 이는 **아다마르 곱(Hadamard product)**과 **서브그래프 유사도(subgraph similarity)**를 사용하는 결합 메커니즘을 채택한다. 마스킹된 노드 주변의 서브그래프를 추출하고, 그 유사도를 계산하며, 아다마르 곱을 적용하여 모든 레이어로부터 다중 그래뉼리티(multi-granularity) 정보를 통합하고 노이즈를 제거한다. 이는 CT 이미지 특유의 노이즈와 흐림 현상을 처리하기 위함이다.

5. 손실 함수:

   • 훈련 목적 함수는 두 가지 손실을 결합한다:
     • 위상 인식 대조 손실(Topology-aware Contrastive Loss): 마스킹된 에지의 양 끝점과 연결된 서브그래프 내의 노드들을 양성 샘м플로, 외부의 노드들을 음성 샘플로 취급한다. 이를 통해 유사한 위상 구조를 서로 가깝게 만든다.
     • 위상 일관성 재구성 손실(Topology-consistence Reconstruction Loss): 0차원 지속성 다이어그램(0-dimensional persistence diagrams)을 사용하여 원본 그래프와 마스킹된 그래프 사이의 위상적 거리를 측정하며, 이를 통해 재구성이 병변의 전역적 위상 구조를 보존하도록 보장한다.

주요 기여

• 에지 중심 마스킹: 그래프 오토인코더에서 노드 대신 에지를 마스크킹하여, 국소 병리적 특징을 포착하는 데 있어 노드 마스킹 방식의 한계를 해결하고 복셀 블록 간의 관계를 명시적으로 모델링하는 방안을 제안하였다.
• 효율적인 아키텍처: GCN 인코더와 특화된 디코더(아다마르 곱 + 서브그래프 유사도)의 통합을 통해, 어텐션 기반 트랜스포머의 이차적 복잡도($O(N^2)$)와 대비되는 선형 복잡도($O(|E|)$)를 구현하여 계산 오버헤드를 크게 줄였다.
• 소규모 데이터셋에서의 효능: 88개의 케이스와 같은 소규모 데이터셋에서도 고성능 이상 탐지가 가능하다는 것을 입증하며, 대규모 사전 학습 데이터셋이나 고성능 GPU 클러스터 없이도 작동함을 보여주었다.
• 새로운 손실 함수: 서브그래프 연결성과 지속적 호몰로지(persistent homology)를 활용하여 병변 특징의 명확성과 경계 정의를 개선하는 결합 위상 인식 대조-재구성 손실을 도입하였다.

실험 결과

본 방법은 LUNA16 데이터셋(888개의 저선량 폐 CT 스캔을 포함하는 LIDC-IDRI의 일부)을 통해 평가되었다.

• 훈련 효율성:

  • Apple M2 (8GB RAM)를 사용한 88개 케이스의 소규모 데이터셋에서, 제안된 방법은 25분의 훈련 시간이 소요되었다.
  • 반면, 베이스라인인 MAE는 30시간, DrasCLR은 35시간이 소요되었다.
  • 888개 케이스 전체 데이터셋에 대해, 제안된 방법은 베이스라인 대비 훈련 시간을 30배 단축시켰으며 계산 비용을 68% 절감했다.

• 성능 (결절 탐지):

  • 제안된 방법은 88개 케이스 서브셋에서 0.779의 **CPM (탐지 점수)**를, 888개 케이스 전체 데이터셋에서 0.822를 달립했다.
  • 이 결과는 다음과 같은 최첨단(SOTA) 베이스라인들과 경쟁할 만한 수준이다:
    • DeepLung (3D CNN): 0.819 (88 케이스) / 0.958 (888 케이스).
    • GraphMAE: 0.755 (88 케이스) / 0.793 (888 케이스).
    • DrasCLR: 0.682 (88 케이스) / 0.762 (888 케이스).
  • 저자는 자신의 방법이 훨씬 적은 자원과 더 작은 훈련 세트를 사용했음에도 불구하고, 베이스라인 중 2위를 기록했으며 단일 단계 결절 탐지 모델과 대등한 성능을 보였다고 언급했다.

• 어블레이션 연구 (Ablation Studies):

  • 마스킹 전략: 에지 마스킹이 마스킹되지 않은 경우와 노드 마스킹 전략 모두보다 우수한 성능을 보였다.
  • 디코더: 결합 디코더(아다마르 곱 + 서브그래프 유사도)가 MLP, 아다마르 곱, 또는 서브그래프 유사도만을 사용한 변형 모델들보다 유의미하게 뛰어난 성능을 보였다.
  • 손실 함수: 결합된 위상 인식 대조-재구성 손실은 각 손실 성분을 단독으로 사용했을 때보다 우수한 결과를 냈다.
  • 데이터 규모: 소규모(88 케이스)와 대규모(888 케이스) 훈련 간의 성능 차이가 미미하여, 제한된 데이터로부터 효과적으로 학습하는 능력을 검증했다.

의의 및 주장

본 논문은 3D CT-GraphEdgeMAE가 자원이 제한된 환경에서 3D CT 이상 탐지를 위한 실행 가능한 솔루션을 제공한다고 주장한다. 어텐션 부하가 큰 패치 기반 트랜스포머에서 그래프 기반 에지 마스킹 방식으로 전환함으로써, 이 방법은 고성능과 대규모 데이터셋 및 고사양 컴퓨팅 하드웨어의 필요성을 성공적으로 분리해 냈다.

저자는 다음과 같이 강조한다:

1. 접근성의 민주화: 소비자급 하드웨어(예: Apple M2)와 소규모 임상 데이터셋에서도 효과적인 훈련을 가능하게 하여, A100/RTX4090 클러스터가 없는 환경에서도 고급 이상 탐지를 사용할 수 있게 한다.
2. 위상 보존: 2D 패치 기반 또는 노드 마스킹 방식에서 손실되기 쉬운 3D 결절의 공간적 및 위상적 관계를 명시적으로 모델링한다.
3. 이차적 피해 감소: 재검사(이에 따른 추가 방사선 노출) 없이도 흐릿하거나 노이즈가 있는 이미지에서 탐지 정확도를 개선함으로써 중요한 임상적 요구를 해결한다.

결론적으로, 본 연구는 위상 인식 그래프 구조와 결합된 소규모 데이터셋 기반의 자기지도 학습이 대규모 사전 학습 패러다임과 대등한 성능을 달성할 수 있음을 보여주며, 효율적인 의료 영상 분석의 진전을 나타낸다.