Interpretable Medical Image Classification using Prototype Learning and Privileged Information

Each language version is independently generated for its own context, not a direct translation.

이 논문은 **"의사들이 왜 그 결정을 내렸는지 설명할 수 있는, 똑똑한 AI"**를 만드는 방법에 대해 이야기합니다.

의료 영상 (예: 폐 CT 스캔) 을 분석할 때, 최신 AI 는 인간 의사보다 더 정확하게 암을 찾아내기도 합니다. 하지만 문제는 "왜 그렇게 판단했는지"를 설명해 주지 않는다는 점입니다. 마치 "이게 암이야"라고만 말하고 이유를 알려주지 않는 의사처럼 말이죠.

이 연구팀은 **"Proto-Caps"**라는 새로운 AI 모델을 개발했는데, 이를 쉽게 설명해 드릴게요.

1. 핵심 아이디어: "비밀 스승"과 "비유 그림"의 만남

이 모델은 두 가지 아주 clever 한 아이디어를 섞었습니다.

비밀 스승 (Privileged Information):
훈련할 때만 사용할 수 있는 '비밀 정보'를 활용합니다. 보통 AI 는 "이게 암인가?"만 배우지만, 이 모델은 훈련 중에는 "이 암 덩어리는 둥글까? (구형성), 가시가 있나? (침상성), 경계가 뚜렷할까? (연결성)"처럼 의사가 중요하게 여기는 세부 특징들도 함께 배웁니다.
- 비유: 시험을 볼 때는 답만 외우는 게 아니라, 문제 풀이 과정과 힌트도 함께 공부하는 학생이라고 생각하세요. 시험 당일에는 힌트가 없어도, 그 과정을 통해 배운 논리로 정답을 맞힙니다.
비유 그림 (Prototype Learning):
AI 가 판단한 이유를 설명할 때, 단순히 숫자만 보여주는 게 아니라 **"이런 모양의 실제 사례"**를 보여줍니다.
- 비유: "이 환자는 '가시가 있는' 종양이니까 암일 확률이 높다"고 말할 때, AI 는 "보시죠, 이 환자와 가장 비슷한 '가시가 있는' 종양 사례 3 개가 여기 있습니다"라고 실제 이미지를 보여주며 설명합니다.

2. Proto-Caps 가 어떻게 작동하나요? (우주선 캡슐 이야기)

이 모델은 **'캡슐 네트워크 (Capsule Network)'**라는 기술을 사용합니다.

캡슐 (Capsule): 일반적인 AI 는 이미지를 전체적으로 뭉개서 보지만, 이 모델은 이미지를 여러 개의 **'작은 우주선 (캡슐)'**으로 나눕니다.
각 우주선은 특정 역할 (예: 둥글기, 가시, 크기 등) 을 담당합니다.
훈련할 때 이 우주선들은 '비밀 스승' (의사의 세부 설명) 을 보고 학습합니다.
그리고 각 우주선은 **"내가 본 이 특징과 가장 비슷한 실제 사례 (비유 그림)"**를 찾아냅니다.

3. 왜 이것이 중요한가요?

기존의 설명 가능한 AI 들은 "이 부분이 중요해서 암이라고 봤어요"라고 **어떤 부분 (히트맵)**을 강조하는 정도였습니다. 하지만 Proto-Caps 는 다음과 같이 더 구체적입니다.

정확도 향상: 단순히 설명만 잘하는 게 아니라, 정확도도 6% 이상 더 높아졌습니다. (암 진단 정확도 93% 달성)
실제 검증: "가시 (Spiculation) 가 있다고 판단했는데, 보여준 '비유 그림'을 보니 가시가 없네?"라고 의사가 바로 확인할 수 있습니다. AI 의 실수를 눈으로 확인하고 수정할 수 있는 것이죠.
적은 정보로도 가능: 의사가 모든 종양에 대해 세부 특징을 다 적어주지 않아도 (훈련 데이터의 10% 만 있어도) AI 는 잘 작동합니다. 이는 실제 병원에서 적용하기 훨씬 수월하다는 뜻입니다.

4. 한 줄 요약

이 연구는 **"AI 가 암을 진단할 때, 단순히 '암이다'라고 말하는 것을 넘어, '이런 특징이 있어서 암이라고 봤고, 비슷한 사례는 여기 있습니다'라고 의사가 눈으로 직접 확인할 수 있게 설명해주는 AI"**를 만들었습니다.

이는 AI 를 단순히 '검은 상자 (Black Box)'가 아니라, 의사와 대화할 수 있는 신뢰할 수 있는 파트너로 만들어주는 중요한 한 걸음입니다.

Each language version is independently generated for its own context, not a direct translation.

논문 개요: Proto-Caps

이 논문은 의료 영상 분류, 특히 폐 결절의 악성도 (malignancy) 예측에서 **해석 가능성 (Interpretability)**과 고성능을 동시에 달성하기 위한 새로운 딥러닝 모델인 Proto-Caps를 제안합니다. 저자들은 학습 과정 중에만 사용 가능한 추가 정보 (Privileged Information) 와 프로토타입 학습 (Prototype Learning) 을 캡슐 네트워크 (Capsule Network) 에 통합하여, 방사선 전문의가 정의한 속성 (attributes) 을 기반으로 한 시각적 추론이 가능한 모델을 개발했습니다.

1. 문제 제기 (Problem)

딥러닝의 블랙박스 문제: 의료 영상 분석에서 딥러닝은 인간 수준의 예측 성능을 보이지만, 그 결정 과정이 불투명하여 임상 현장에서 신뢰를 얻기 어렵습니다.
기존 해석 가능 모델의 한계:
- Privileged Information (학습 시만 가능한 정보) 활용: 기존 연구 (예: Shen et al., LaLonde et al.) 는 학습 시 방사선 의사가 표시한 결절의 특성 (구형도, 경계, 침습 등) 을 활용하여 성능을 높였으나, 예측된 속성 점수가 틀릴 경우 이를 검증할 방법이 부족했습니다.
- 프로토타입 기반 모델의 한계: 기존 프로토타입 네트워크는 "어떤 학습 사례와 유사한가"는 보여줄 수 있지만, "왜 그런 특징이 중요한가"에 대한 구체적인 속성 기반 설명을 제공하지 못했습니다.
목표: 성능 저하 없이 해석 가능성을 내재화 (Explainable-by-design) 하고, 예측된 속성 값의 신뢰성을 시각적으로 검증할 수 있는 모델 구축.

2. 방법론 (Methodology)

제안된 Proto-Caps는 캡슐 네트워크 아키텍처를 기반으로 하며, 다음과 같은 핵심 요소들로 구성됩니다.

가. 아키텍처 구조

백본 (Backbone): 2D 합성곱 층과 3 개의 캡슐 레이어로 구성됩니다.
- 입력: $1 \times 32 \times 32$ 크기의 영상.
- Primary Capsules: 저수준 특징을 8 개의 캡슐로 분리.
- Dense Capsules: 각 캡슐이 사전 정의된 8 가지 속성 (Subtlety, Internal Structure, Calcification, Sphericity, Margin, Lobulation, Spiculation, Texture) 을 나타내며, 16 차원 벡터로 인코딩됩니다.
다중 헤드 (Multi-head) 학습:
1. Target Head: 캡슐 인코딩을 결합하여 최종 악성도 (Malignancy) 를 예측합니다. (KL 발산을 손실 함수로 사용하여 의사의 불확실성을 반영).
2. Reconstruction Head: 노듈의 분할 마스크 (Segmentation Mask) 를 재구성하여 모델이 올바른 특징을 학습하도록 돕습니다.
3. Attribute Head: 각 캡슐 벡터를 별도의 선형 계층을 통해 통과시켜 해당 속성의 점수를 예측합니다.

나. 프로토타입 학습 및 검증 (Prototype Learning)

프로토타입 생성: 각 속성 클래스당 2 개의 프로토타입을 학습합니다 (총 8~12 개). 학습 데이터의 대표적 부분집합으로, 실제 원본 이미지가 저장되어 시각화됩니다.
손실 함수 (Loss Functions):
- Clustering Loss ( $L_{clu}$ ): 샘플의 캡슐 벡터를 올바른 속성 클래스의 프로토타입에 가깝게 묶습니다.
- Separation Loss ( $L_{sep}$ ): 잘못된 속성 클래스의 프로토타입으로부터 거리를 둡니다.
추론 (Inference) 단계: 예측된 속성 값은 학습된 밀집 계층 (dense layers) 의 출력이 아니라, **가장 가까운 프로토타입의 실제 속성 값 (Ground Truth)**으로 대체됩니다. 이를 통해 모델이 "이 입력은 이 프로토타입과 유사하므로 이 속성을 가진다"는 논리를 제공합니다.

다. 전체 손실 함수

$L = L_{mal} + \lambda_{recon} \cdot L_{recon} + L_{attr} + 0.125 \cdot (L_{clu} + 0.1 \cdot L_{sep})$
여기서 $\lambda_{recon}$ 은 재구성 손실의 가중치이며, 프로토타입 관련 가중치는 실험적으로 결정되었습니다.

3. 주요 기여 (Key Contributions)

최초의 통합 접근법: 의료 분류 작업에서 Privileged Information과 Prototype Learning을 결합하여 해석력을 극대화한 최초의 방법론을 제시했습니다.
새로운 아키텍처 (Proto-Caps): 캡슐 네트워크의 장점 (속성별 인코딩) 과 프로토타입 학습 (시각적 검증) 을 결합한 네트워크 구조를 제안했습니다.
성능 및 해석성 동시 달성: LIDC-IDRI 데이터셋에서 기존 설명 가능한 모델 및 비설명 가능한 모델보다 우수한 성능을 달성하면서도, 예측의 근거를 시각적으로 제시합니다.

4. 실험 결과 (Results)

데이터셋: LIDC-IDRI (1,018 개의 흉부 CT 스캔, 3mm 이상 폐 결절).
성능 비교 (Within-1-Accuracy):
- 악성도 예측: 제안된 방법은 **93.0%**의 정확도를 기록하여, 기존 설명 가능한 모델 (X-Caps: 90.4%) 보다 약 6% 이상 높은 성능을 보였습니다.
- 속성 예측: 8 가지 속성 모두에서 기존 방법 (HSCNN, CapsNet 등) 보다 높은 정확도를 달성했습니다 (예: Sphericity 95.4%, Spiculation 93.3%).
데이터 부족 실험 (Data Reduction):
- 학습 데이터 중 속성 라벨이 있는 샘플을 10% 로 줄여도 악성도 예측 성능 (92.4%) 은 유지되었습니다.
- 속성 라벨이 아예 없는 경우 (0%) 도 유사한 성능을 보였으나, 해석 가능성은 사라졌습니다. 이는 모델이 속성 정보를 명시적으로 학습하지 않아도 성능을 낼 수 있음을 시사하지만, 해석을 위해서는 라벨이 필요함을 보여줍니다.
정성적 분석:
- 프로토타입과 입력 이미지의 불일치를 통해 잘못된 속성 예측을 식별할 수 있었습니다.
- 속성 예측의 정확도와 악성도 예측의 정확도 간에 강한 상관관계가 있음을 확인했습니다.

5. 의의 및 결론 (Significance)

임상적 신뢰성 향상: 단순히 "악성이다/양성이다"라고 예측하는 것을 넘어, 어떤 특징 (예: 경계 불명확, 침습성 등) 을 근거로 그 결론에 도달했는지 시각적 프로토타입을 통해 방사선 전문의가 검증할 수 있게 합니다.
효율적인 학습: 모든 학습 데이터에 상세한 속성 라벨이 없더라도 (10% 만 사용 시) 고성능을 유지하므로, 라벨링 비용이 큰 의료 데이터셋에서도 적용 가능성이 높습니다.
미래 연구 방향: 의료 보고서와 같은 다른 형태의 Privileged Information 활용이나, 3D 기반 캡슐 네트워크 확장 등을 통해 고위험 의료 분야에서의 적용 범위를 넓힐 수 있음을 시사합니다.

이 연구는 **"설명 가능성은 성능의 희생이 아니라, 설계 단계에 내재화될 때 오히려 성능을 향상시킬 수 있다"**는 가설을 입증하며, 의료 AI 의 신뢰성 있는 도입을 위한 중요한 발걸음이 되었습니다.