Effect of Patch Size on Fine-Tuning Vision Transformers in Two-Dimensional and Three-Dimensional Medical Image Classification

이 논문은 2D 및 3D 의료 영상 분류 작업에서 비전 트랜스포머 (ViT) 의 패치 크기가 작을수록 (1, 2, 4) 계산 비용 증가에도 불구하고 분류 성능이 크게 향상되며, 특히 패치 크기 1, 2, 4 모델을 앙상블하면 성능이 더욱 개선됨을 12 개 의료 영상 데이터셋을 통해 실증했습니다.

핵심

이 논문은 **"의사들이 엑스레이나 MRI 같은 의료 영상을 볼 때, 컴퓨터가 그 이미지를 어떻게 잘게 쪼개서 분석해야 가장 정확하게 병을 찾아낼 수 있을까?"**라는 아주 중요한 질문을 던집니다.

비유를 들어 쉽게 설명해 드릴게요.

🧩 핵심 비유: 퍼즐 조각의 크기 (Patch Size)

생각해 보세요. 거대한 의료 영상을 퍼즐이라고 가정해 봅시다.
이전까지 컴퓨터 비전 (Vision Transformer) 기술은 이 퍼즐을 **너무 큰 조각 (예: 14x14 픽셀)**으로만 잘라서 분석했습니다. 마치 거대한 조각으로만 퍼즐을 맞추려고 하니까, 디테일한 부분 (예: 종양의 미세한 가장자리, 혈관의 작은 변화) 을 놓치고 전체적인 모양만 대충 파악하는 셈이죠.

이 연구는 **"조각을 더 작게 (1x1, 2x2, 4x4 등) 자르면 어떨까?"**를 실험했습니다.

📝 이 연구가 발견한 놀라운 사실들

1. 작을수록 더 똑똑해집니다 (2D & 3D 모두)

• 2D 이미지 (엑스레이, 피부 사진 등): 퍼즐 조각을 아주 작게 자를수록 (작은 Patch Size), 컴퓨터가 병을 찾는 정확도가 약 12% 이상이나 올라갔습니다.
• 3D 이미지 (CT, MRI 등): 이건 더 극적이었습니다. 3D 입체 영상을 작은 조각으로 쪼개면 정확도가 약 24%까지 급상승했습니다.
• 왜 그럴까요? 작은 조각은 마치 고배율 현미경처럼 작동합니다. 병의 아주 미세한 특징까지 놓치지 않고 포착하기 때문이죠. 큰 조각으로는 놓치던 중요한 단서들을 작은 조각들이 모두 찾아냅니다.

2. 하지만 '비싼' 가격표가 붙어 있습니다

• 퍼즐 조각을 작게 자르면 조각의 개수가 기하급수적으로 늘어납니다.
• 비유: 1 개의 큰 조각을 분석하는 데 1 초 걸린다면, 그걸 100 개의 작은 조각으로 나누어 분석하려면 컴퓨터는 훨씬 더 많은 일을 해야 합니다.
• 결과: 정확도는 엄청나게 좋아졌지만, 컴퓨터가 계산하는 데 드는 **시간과 전기세 (컴퓨팅 비용)**도 함께 폭증했습니다. 특히 3D 영상은 조각이 너무 많아져서 계산량이 64 배까지 불어날 수도 있습니다.

3. "모두 합치면 더 완벽해" (앙상블 전략)

• 연구진은 "작은 조각 (1), 조금 큰 조각 (2), 중간 크기 조각 (4) 으로 각각 분석한 뒤, 그 결과를 모두 섞어서 (Ensemble) 최종 판단을 내리면" 가장 좋은 결과를 얻는다는 것을 발견했습니다.
• 비유: 한 명의 전문가가 모든 것을 보는 것보다, 세 명의 전문가가 각각 다른 관점 (미세한 것, 중간, 전체) 으로 보고 의견을 합치는 것이 가장 정확한 진단을 내린다는 뜻입니다.

🏥 이 연구가 왜 중요한가요?

1. 의료 AI 의 새로운 기준: 그동안 의료 영상 분석 AI 는 "큰 조각"을 쓰는 것이 당연시되었는데, 이 연구는 **"작은 조각을 써야 더 정확하다"**는 것을 증명했습니다.
2. 단일 컴퓨터로도 가능: 보통 이런 거대한 실험을 하려면 슈퍼컴퓨터나 수백 개의 그래픽 카드가 필요하다고 알려졌는데, 이 연구는 일반적인 고성능 컴퓨터 한 대로도 충분히 정밀한 분석이 가능하다는 것을 보여줬습니다.
3. 실제 임상 적용의 고민: 정확도는 좋아졌지만, 계산 비용이 너무 많이 들어 병원 현장에서 실시간으로 쓰기엔 무거울 수 있습니다. 그래서 **"정확도와 속도 사이의 균형"**을 찾는 것이 앞으로의 과제가 되었습니다.

💡 한 줄 요약

"의료 영상을 분석할 때, 퍼즐 조각을 아주 작게 잘게 쪼개고 여러 관점에서 합쳐서 보면, 컴퓨터가 병을 훨씬 더 정확하게 찾아낼 수 있습니다. 다만, 그 대가로 컴퓨터가 더 많이 일해야 한다는 점은 잊지 말아야 합니다."

이 연구는 의료 AI 가 더 정밀한 진단을 내릴 수 있는 새로운 길을 제시했지만, 동시에 그 길을 걷기 위해 필요한 '에너지'에 대한 현실적인 고려도 필요함을 일깨워줍니다.

기술 요약

1. 연구 배경 및 문제 제기 (Problem)

• 비전 트랜스포머 (ViT) 의 중요성: ViT 와 그 변형 모델들은 컴퓨터 비전 분야에서 최첨단 (SOTA) 성능을 보이며, 대규모 비전 및 비전 - 언어 기반 모델의 백본으로 널리 사용되고 있습니다.
• 미세한 설계 선택의 간과: ViT 아키텍처의 개선에 대한 연구는 활발하지만, ViT 파이프라인의 초기 설계 단계에서 가장 중요한 요소 중 하나인 '패치 크기 (Patch Size)' 의 영향력은 의료 영상 분야에서 충분히 탐구되지 않았습니다.
• 의료 영상의 특수성: 의료 영상은 자연 이미지와 달리 2 차원 (2D) 과 3 차원 (3D) 형식이 혼재되어 있으며, 기존 연구들은 주로 자연 이미지나 2D 데이터에 집중했습니다. 또한, 대부분의 연구가 대규모 컴퓨팅 자원 (다수 GPU/TPU 클러스터) 을 사용한 '처음부터 학습 (Training from scratch)' 방식에 의존하여, 소규모 의료 데이터셋에서 주로 사용되는 '미세 조정 (Fine-tuning)' 접근법과의 괴리가 존재했습니다.
• 핵심 질문: 의료 영상 분류에서 패치 크기를 어떻게 설정해야 최적의 성능을 낼 수 있으며, 2D 와 3D 영상에서 이 영향은 어떻게 달라지는가?

2. 방법론 (Methodology)

• 데이터셋: MedMNIST V2 컬렉션에서 다양한 영상 모달리티를 포함한 12 개의 의료 영상 데이터셋을 사용했습니다.
  • 2D 데이터셋 (7 개): Blood, Breast, Derma, OCT, OrganS, Pneumonia, Retina 등.
  • 3D 데이터셋 (5 개): Adrenal, Fracture, Nodule, Synapse, Vessel 등 (CT, MRI 등).
  • 입력 크기: 실험의 계산 비용 (GPU 메모리) 을 단일 GPU 에서 처리 가능하도록 모든 이미지를 28x28 (2D) 또는 28x28x28 (3D) 크기로 통일했습니다.
• 모델 및 설정:
  • 모델: ImageNet 사전 학습 가중치를 기반으로 한 ViT-Small 모델을 사용했습니다.
  • 3D 처리: 2D 사전 학습 가중치를 3D 커널로 확장하는 'Weight Inflation' 기법과 3D 공간 보간법을 적용하여 3D 데이터에 적용했습니다.
  • 패치 크기 실험: 입력 해상도 (28) 의 약수인 1, 2, 4, 7, 14, 28 크기의 패치를 사용하여 미세 조정 (Fine-tuning) 을 수행했습니다.
  • 학습 환경: 단일 NVIDIA RTX 4090 GPU 환경에서 실행하여 실용성을 검증했습니다.
• 앙상블 전략: 패치 크기 1, 2, 4 로 학습된 모델들의 예측을 단순 평균 (Averaging) 하여 융합하는 앙상블 전략을 적용하여 추가 성능 향상을 검증했습니다.

3. 주요 기여 (Key Contributions)

1. 체계적인 패치 크기 영향 분석: 2D 및 3D 의료 영상 분류에서 ViT 의 패치 크기가 성능에 미치는 영향을 포괄적으로 평가한 최초의 연구입니다.
2. 실용적인 실험 환경 증명: 고가의 대규모 컴퓨팅 인프라 없이도, 적절한 데이터셋 크기를 선택하면 단일 GPU에서 상세한 토큰화 (Tokenization) 연구가 가능함을 입증했습니다.
3. 의료 도메인 특화 통찰: 자연 이미지가 아닌 2D/3D 의료 영상에 특화된 패치 크기 선택 가이드라인을 제시했습니다.
4. 오픈 소스 공개: 모든 코드와 구현을 GitHub 에서 공개하여 연구의 투명성과 재현성을 보장했습니다.

4. 실험 결과 (Results)

• 패치 크기와 성능의 상관관계:
  • 일반적인 경향: 패치 크기가 작아질수록 (1, 2, 4) 분류 성능이 일관되게 향상되었습니다. 특히 패치 크기 1, 2, 4가 거의 모든 데이터셋에서 최상의 결과를 보였습니다.
  • 2D 데이터셋: 패치 크기 28(전체 이미지) 대비 패치 크기 2 사용 시 균형 정확도 (Balanced Accuracy) 가 최대 12.78% 까지 향상되었습니다.
  • 3D 데이터셋: 3D 데이터에서 패치 크기의 영향이 더욱 극명했습니다. 패치 크기 14 대비 패치 크기 1 사용 시 균형 정확도가 최대 23.78% 까지 향상되었습니다.
  • 성능 저하 요인: 큰 패치 크기 (14, 28) 는 국소적인 세부 정보 (Local details) 와 질감 정보를 포착하지 못해 진단적으로 중요한 특징을 놓치는 경향이 있었습니다.
• 앙상블 효과: 패치 크기 1, 2, 4 로 학습된 모델의 예측을 융합한 앙상블 모델이 대부분의 경우 개별 모델보다 더 높은 성능을 기록했습니다.
• 주의 시점 (Small-scale 데이터): BreastMNIST 와 RetinaMNIST 와 같이 학습 데이터가 매우 적은 소규모 데이터셋에서는 패치 크기에 따른 성능 패턴이 불명확하거나 일관되지 않았습니다.
• 계산 비용 (Trade-off):
  • 성능 향상은 계산 비용 증가를 동반합니다.
  • 2D: 패치 크기를 28 에서 1 로 줄이면 GFLOPs 가 약 0.04 에서 16.71 로 약 4 배 증가합니다.
  • 3D: 패치 크기를 28 에서 1 로 줄이면 GFLOPs 가 약 0.40 에서 800 이상으로 약 5~7 배 (또는 그 이상) 급증합니다. 이는 3D ViT 에서 토큰 수가 입방적으로 증가하기 때문입니다.
• 어텐션 맵 분석: 작은 패치 크기 (P2) 를 사용한 모델은 진단적으로 중요한 영역에 더 집중적이고 세부적인 어텐션 (Attention) 을 보이는 반면, 큰 패치 (P28) 는 균일하고 덜 정보적인 어텐션 분포를 보여주었습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 의료 AI 설계 가이드라인: 의료 영상 분석을 위한 ViT 기반 파이프라인 설계 시, **작은 패치 크기 (특히 1~4)**를 사용하는 것이 성능 향상에 결정적임을 입증했습니다. 이는 미세한 병변이나 구조적 특징을 포착하는 데 필수적입니다.
• 자원 효율성: 대규모 컴퓨팅 자원이 없더라도 단일 GPU 환경에서 패치 크기 최적화 연구가 가능함을 보여주어, 많은 연구 그룹이 접근 가능한 실용적인 연구 방향을 제시했습니다.
• 한계 및 향후 과제:
  • 3D 데이터의 경우 성능 향상 폭이 크지만 계산 비용이 매우 높아 실시간 임상 적용에 제약이 있을 수 있습니다.
  • 실험이 MedMNIST (저해상도) 에 기반했으므로, 실제 고해상도 임상 데이터에서도 동일한 경향이 유지되는지 추가 검증이 필요합니다.
  • ViT-Small 만 사용했으므로, ViT-Large 나 Huge 와 같은 더 큰 모델에서도 동일한 경향이 나타날지 연구가 필요합니다.

요약하자면, 이 연구는 의료 영상 분류에서 ViT 의 패치 크기를 줄이는 것이 계산 비용 증가라는 대가를 치르더라도 성능 향상을 가져오는 핵심 요소임을 2D 와 3D 데이터 모두에서 체계적으로 입증했습니다.