VariViT: A Vision Transformer for Variable Image Sizes

1. 문제 정의 (Problem)

비전 트랜스포머 (ViT) 는 자기 주의 (self-attention) 메커니즘을 통해 이미지 표현 학습에서 최첨단 성능을 보이지만, 의료 영상 분석에는 다음과 같은 근본적인 한계가 존재합니다.

• 고정된 입력 크기 제약: 기존 ViT 는 이미지를 고정된 크기의 패치 (patch) 로 분할하므로, 입력 이미지를 미리 정해진 크기로 리사이징 (resizing), 패딩 (padding), 또는 크롭 (cropping) 해야 합니다.
• 의료 영상의 비정형성: 종양과 같은 병변은 모양과 크기가 매우 불규칙합니다. 고정된 크기의 바운딩 박스로 크롭할 경우, 전경 (종양) 대 배경 (정상 조직) 의 비율이 이미지마다 크게 달라집니다.
• 정보 손실 및 아티팩트:
  • 리사이징: 의료 영상을 고정 크기로 리사이징하면 중요한 미세한 정보가 손실되거나 인위적인 아티팩트가 생성되어 진단에 악영향을 줄 수 있습니다.
  • 배경 노이즈: 작은 병변을 포함하는 큰 크롭의 경우, 모델이 관심 영역 (ROI) 이 아닌 배경에 과도하게 주의를 기울여 (attention) 계산 자원을 낭비합니다.
• 계산 효율성: 큰 이미지를 처리하는 것은 계산 비용이 많이 들지만, 작은 이미지로 무작위 크롭하는 것은 정보 손실의 위험이 있습니다.

2. 제안 방법론 (Methodology)

저자들은 VariViT라는 새로운 아키텍처를 제안하여 가변적인 이미지 크기를 처리하면서도 패치 크기는 일정하게 유지하는 방법을 개발했습니다.

가. Center and Select (위치 임베딩 조정 기법)

기존 ViT 는 고정된 위치 임베딩 (positional embedding) 을 사용하거나 고해상도에서 선형 보간 (interpolation) 을 적용합니다. 하지만 3D 의료 영상에서 보간은 정보 손실과 계산 복잡도를 증가시킵니다.

• 핵심 아이디어: 3D 종양 크롭은 종양 질량 중심 (center of mass) 을 기준으로 정렬되어 있다는 점을 활용합니다.
• 구현:
  1. 가장 큰 이미지 크기에 해당하는 고정된 3D 정현파 (sinusoidal) 위치 임베딩 그리드를 초기화합니다.
  2. 입력된 이미지의 크기가 작아지면, 이 큰 그리드의 **중앙 (center)**을 기준으로 해당 크기에 맞는 서브셋을 **선택 (select)**합니다.
  3. 보간 (interpolation) 을 사용하지 않고 원본 위치 정보를 그대로 추출하므로 계산 비용이 낮고 정보 손실이 없습니다.

나. 배칭 전략 (Batching Strategies)

서로 다른 크기의 이미지를 배치 (batch) 로 처리하는 것은 기존 프레임워크에서 어렵습니다. 이를 해결하기 위해 두 가지 전략을 도입했습니다.

1. Custom Batch Sampler (CBS): 같은 크기의 이미지들을 하나의 배치로 묶습니다. 배치 간에는 크기가 달라질 수 있지만, 배치 내에서는 일관성을 유지하여 효율성을 높입니다.
2. Gradient Accumulation (GA): 미니배치 크기를 1 로 설정하고 여러 번의 전진/역전파를 통해 그래디언트를 누적한 후 가중치를 업데이트합니다. 이를 통해 서로 다른 크기의 이미지를 하나의 배치로 처리할 수 있으며, 메모리 효율성을 극대화합니다.

3. 주요 기여 (Key Contributions)

1. 유연한 위치 임베딩 전략: 다양한 이미지 크기에 맞춰 중앙 정렬을 기반으로 한 'Center and Select' 방식을 도입하여, 보간 없이도 정확한 위치 정보를 유지합니다.
2. 계산 효율성 향상: 가변 크기의 이미지를 효율적으로 처리하는 새로운 배칭 전략 (CBS, GA) 을 제안하여 학습 및 추론 시간을 단축합니다.
3. 의료 영상 적용 검증: 뇌 MRI 데이터를 이용한 글리오마 유전자형 예측 (IDH mutation status) 및 뇌종양 분류 (원발성 vs 전이성) 작업에서 기존 모델 대비 우수한 성능을 입증했습니다.

4. 실험 결과 (Results)

두 가지 3D 뇌 MRI 데이터셋 (글리오마 데이터셋, 뇌종양 데이터셋) 에서 VariViT 를 평가했습니다.

• 성능 비교 (IDH 변이 상태 예측):
  • VariViT-GA는 AUC 0.942, F1-Score 0.755, MCC 0.709 를 기록하여 고정 크기 기반의 ResNet-18, Vanilla ViT, 그리고 Pix2Struct(가변 크기 모델) 보다 우수한 성능을 보였습니다.
  • VariViT-CBS 역시 기존 모델들을 상회하거나 동급의 성능을 내면서 학습 시간을 단축했습니다.
• 성능 비교 (뇌종양 분류):
  • VariViT-GA 는 F1-Score 0.763을 기록하여 Vanilla ViT(0.752) 및 기타 베이스라인 모델보다 높은 정확도를 보였습니다.
• 학습 효율성:
  • 제안된 배칭 전략을 통해 기존 아키텍처 대비 최대 30% 까지 계산 시간을 단축했습니다.
  • t-SNE 시각화 결과, VariViT 는 종양 영역에 더 명확하게 집중하여 클래스 간 분리가 더 잘 됨을 확인했습니다.
• Ablation Study:
  • 다양한 위치 임베딩 전략 (독립 고정, 보간 고정, 보간 학습) 을 비교한 결과, 제안한 Center and Select 방식이 가장 우수한 성능을 보였습니다.

5. 의의 및 결론 (Significance)

• 의료 영상 분석의 패러다임 전환: 의료 영상에서 발생하는 불규칙한 병변 크기와 배경 비율의 문제를 해결하기 위해, 리사이징이나 패딩 없이 원본 정보를 최대한 보존하면서 ViT 를 적용할 수 있는 길을 열었습니다.
• 효율성과 정확성의 균형: 계산 비용이 많이 드는 리사이징이나 복잡한 토큰 드롭 (token dropping) 기법 없이, 단순하고 효율적인 '중앙 선택' 및 '배칭' 전략으로 성능과 속도를 동시에 개선했습니다.
• 확장성: 뇌 MRI 에 국한되지 않고, 다양한 의료 영상 모달리티와 고해상도 이미지 처리에도 적용 가능한 유연한 프레임워크를 제시했습니다.

결론적으로, VariViT 는 의료 영상 분석에서 ViT 의 잠재력을 극대화하면서도 계산적 비효율성과 정보 손실 문제를 해결한 실용적인 솔루션입니다.