Unified Medical Image Tokenizer for Autoregressive Synthesis and Understanding

논문 요약: Unified Medical Image Tokenizer for Autoregressive Synthesis and Understanding (MedITok)

1. 문제 정의 (Problem)

의료 영상 분야에서 생성형 AI 와 이해 모델의 발전을 위한 핵심 요소인 **이미지 토크나이저 (Image Tokenizer)**는 여전히 다음과 같은 한계에 직면해 있습니다.

• 구조와 의미의 분리: 기존 토크나이저는 두 가지 중 하나에 치중하는 경향이 있습니다.
  • 합성 중심 (Synthesis-oriented): VQGAN 등 픽셀 수준의 재구성에 최적화되어 저수준 구조를 잘 보존하지만, 임상적 의미 (고수준 개념) 를 인코딩하지 못합니다.
  • 이해 중심 (Understanding-driven): CLIP 등 텍스트 의미 정렬에 최적화되어 고수준 개념을 잘 파악하지만, 공간 구조와 질감 (texture) 을 왜곡하거나 잃어버립니다.
• 의료 도메인의 특수성: 의료 임무는 정밀한 해부학적 구조 (저수준) 와 임상적 의미 (고수준) 가 모두 필수적입니다. 그러나 기존 접근법은 이 두 가지를 동시에 학습하려 할 때 **상호 간섭 (Gradient Interference)**이 발생하여 성능이 저하됩니다.
• 데이터 부족: 대규모 의료 이미지 - 텍스트 쌍 (Paired Data) 은 귀하며, 대부분의 의료 이미지는 레이블이 없는 Unpaired 데이터로 존재합니다. 기존 방법들은 대규모 쌍 데이터에 의존하여 이 귀한 데이터를 활용하지 못합니다.

2. 방법론 (Methodology)

저자들은 이러한 문제를 해결하기 위해 MedITok이라는 통합 의료 이미지 토크나이저를 제안하며, 이를 훈련시키기 위해 2 단계 훈련 프레임워크를 도입했습니다.

• 2 단계 훈련 프레임워크:

  1. 시각 표현 정렬 단계 (Visual Representation Alignment, Stage 1):
     • 목적: 대규모 **레이블 없는 의료 이미지 (Unpaired Images)**를 활용하여 재구성 충실도 (Reconstruction Fidelity) 를 확보하고 기초적인 의미 인식을 확립합니다.
     • 방식: 사전 훈련된 비전 인코더 (BiomedCLIP Vision Encoder) 를 사용하여 잠재 공간 (Latent Space) 에 가벼운 의미 제약 (Contrastive Loss) 을 가합니다. 이는 재구성 목표와 의미 목표 간의 직접적인 충돌을 완화합니다.
  2. 텍스트 의미 정렬 단계 (Textual Semantic Alignment, Stage 2):
     • 목적: **이미지 - 캡션 쌍 (Image-Text Pairs)**을 사용하여 세밀한 임상 의미 (Fine-grained Clinical Semantics) 를 주입합니다.
     • 방식: 1 단계에서 학습된 토크나이저를 미세 조정 (Fine-tuning) 하여 이미지 토큰을 상세한 임상 캡션의 텍스트 임베딩과 정렬합니다.

• 데이터 규모:

  • Stage 1: 3,340 만 개 이상의 의료 이미지 (9 가지 모달리티 포함).
  • Stage 2: 240 만 개 이상의 고품질 의료 이미지 - 캡션 쌍.
  • 전처리: 저해상도, 정보 부족, 비임상적 이미지 (차트, 표 등) 를 제거하는 엄격한 품질 관리 (Quality Control) 프로세스를 적용했습니다.

• 아키텍처:

  • ViT 기반의 인코더/디코더와 32,768 개의 코드북 벡터를 가진 벡터 양자화 (Vector Quantization, VQ) 모듈을 사용합니다.
  • 이산 토큰 (Discrete Tokens) 을 생성하여 자동회귀 (Autoregressive, AR) 모델과의 호환성을 보장합니다.

3. 주요 기여 (Key Contributions)

1. 새로운 훈련 프레임워크 제안: 재구성과 의미 학습 간의 간섭을 해결하고, 레이블 없는 대규모 의료 데이터를 효과적으로 활용하는 2 단계 훈련 방식을 제시했습니다.
2. MedITok 개발: 저수준 구조 정보와 고수준 임상 의미를 통합적으로 인코딩하는 최초의 통합 의료 이미지 토크나이저입니다.
3. 광범위한 벤치마크 성능: 9 가지 영상 모달리티 (CT, MRI, X-ray, 현미경 등) 와 4 가지 작업 군 (인코딩, 합성, 이해 등) 에 걸쳐 30 개 이상의 데이터셋에서 SOTA(State-of-the-Art) 성능을 입증했습니다.
4. 오픈 소스 및 데이터 접근성: 모델, 코드, 그리고 라이선스를 준수한 데이터 접근 방법을 공개하여 의료 AI 커뮤니티의 발전에 기여합니다.

4. 실험 결과 (Results)

MedITok 은 다양한 태스크에서 기존 모델 (VQGAN, UniTok, MedVAE, PUMIT 등) 을 압도하는 성능을 보였습니다.

• 저수준 인코딩 (이미지 재구성):
  • 8 가지 모달리티에서 rFID, PSNR, SSIM 지표에서 가장 높은 재구성 충실도를 달성했습니다.
  • 특히 16 배 다운샘플링 (16x) 을 사용했음에도 불구하고, 8 배 다운샘플링을 사용하는 다른 모델들보다 우수한 재구성을 보여주며 압축 효율성이 뛰어납니다.
  • 재구성된 이미지에서 진단에 필요한 세부 정보 (색상, 병변 형태 등) 가 잘 보존됨을 분류 작업 (mAP, AUC) 을 통해 입증했습니다.
• 고수준 인코딩 (이미지 분류):
  • 선형 프로빙 (Linear Probing) 실험에서 피부 병변, 안저, 병리, 초음파, X-ray 등 다양한 모달리티의 분류 작업에서 가장 높은 정확도를 기록했습니다.
  • 이는 MedITok 이 단순한 픽셀 인코딩을 넘어 임상적으로 유의미한 고수준 의미를 잘 학습했음을 의미합니다.
• 의료 이미지 합성 (Generative Tasks):
  • LlamaGen 프레임워크에 통합하여 모달리티 조건부 및 피부 질환 조건부 이미지 합성 실험을 수행했습니다.
  • 생성된 이미지의 품질 (gFID) 이 기존 모델보다 우수하며, 시각적 튜링 테스트 (Visual Turing Test) 에서 전문의가 실제 이미지와 합성 이미지를 구분하기 어려운 수준을 보여주었습니다.
• 의료 이미지 이해 (VQA):
  • LLaVA-Med 프레임워크에 통합하여 시각적 질문 답변 (VQA) 태스크에서 VQA-RAD 와 SLAKE 벤치마크에서 최고 성능을 달성했습니다.
  • 다른 모델들이 놓치거나 잘못 해석하는 임상적 세부 사항 (예: "양측성 흉수" vs "우측 흉수") 을 정확하게 식별하고 설명하는 능력을 입증했습니다.

5. 의의 및 결론 (Significance)

• 의료 멀티모달 AI 의 기반 (Foundation): MedITok 은 의료 영상 합성 (Synthesis) 과 이해 (Understanding) 를 하나의 통합된 잠재 공간에서 가능하게 하여, 차세대 의료용 자동회귀 (AR) 모델 개발의 핵심 구성 요소가 됩니다.
• 데이터 효율성: 레이블이 없는 대규모 의료 데이터를 효과적으로 활용할 수 있는 방법을 제시함으로써, 의료 분야에서 쌍 데이터 (Paired Data) 부족이라는 근본적인 문제를 해결하는 실마리를 제공합니다.
• 확장성: 현재 2D 이미지에 최적화되어 있지만, 3D 볼륨 데이터나 다양한 다운스트림 작업으로 확장 가능한 유연한 아키텍처를 갖추고 있습니다.

이 연구는 의료 AI 가 단순한 분류나 검출을 넘어, 복잡한 임상 문맥을 이해하고 현실적인 의료 영상을 생성할 수 있는 강력한 기반을 마련했다는 점에서 의의가 큽니다.