Beyond Attention Heatmaps: How to Get Better Explanations for Multiple Instance Learning Models in Histopathology

이 논문은 디지털 병리학에서 다중 인스턴스 학습 (MIL) 모델의 주석 열지도 (heatmap) 의 유효성을 검증하기 위한 새로운 평가 프레임워크를 제시하고, 기존 주의 메커니즘 기반 방법보다 교란 (perturbation) 및 계층별 관련성 전파 (LRP) 같은 기법이 모델 의사결정 메커니즘을 더 정확하게 반영하여 생물학적 통찰력을 제공할 수 있음을 대규모 벤치마크 실험을 통해 입증했습니다.

핵심

1. 배경: 거대한 퍼즐과 AI 의 추리

병리학자들은 거대한 조직 슬라이드 (Whole Slide Image) 를 현미경으로 보며 병을 진단합니다. 이 슬라이드는 수천 개의 작은 조각 (패치) 으로 나뉘어 있습니다.

• AI 의 역할: AI 는 이 수천 개의 조각을 한 번에 보고 "이 환자는 암입니다"라고 결론을 내립니다. 이를 **다중 인스턴스 학습 (MIL)**이라고 합니다.
• 기존의 문제: AI 가 "어떤 조각이 가장 중요해서 암이라고 판단했는지" 보여주기 위해 **히트맵 (Heatmap, 열지도)**을 그립니다. 마치 범인이 숨어있는 곳을 빨간색으로 표시하는 것 같죠.
  • 하지만 지금까지는 AI 가 만든 이 '빨간색 지도'가 진짜로 AI 가 생각한 이유와 일치하는지 확인해 본 적이 거의 없었습니다. 그냥 "오, 빨간색이 여기 있네? 아마 여기가 중요하겠지?"라고 믿고 넘어갔던 거죠.

2. 연구의 핵심: "지도가 진짜인가, 가짜인가?"를 검증하는 방법

저자들은 이 지도가 AI 의 진짜 추리 과정을 faithfully(충실하게) 반영하는지 확인하는 새로운 방법을 개발했습니다.

• 비유: "범인 제거 게임"
  • AI 가 "이 빨간색 조각들이 중요해서 암이라고 판단했다"고 한다면, 그 빨간색 조각들을 하나씩 지워보면 AI 의 결론이 어떻게 변할까요?
  • 진짜 지도라면: 중요한 조각을 지울수록 AI 는 "아, 이건 아니야"라고 결론을 바꾸거나 확신이 떨어질 것입니다.
  • 가짜 지도라면: 중요한 조각을 지워도 AI 는 여전히 "암이야!"라고 고집을 부릴 것입니다. (즉, 지도가 AI 의 실제 생각과 다르다는 뜻이죠.)
  • 저자들은 이 **'조각 지우기 실험 (Patch Flipping)'**을 통해 각 설명 방법의 신뢰도를 점수화했습니다.

3. 실험 결과: "누가 진짜 범인 (설명 방법) 인가?"

저자들은 6 가지 다른 설명 방법 (AI 가 어떻게 생각했는지 보여주는 도구) 을 비교했습니다. 결과는 놀라웠습니다.

• 실패한 방법 (Attention Heatmap):
  • 가장 흔히 쓰이는 '주의 (Attention)' 지도는 가장 신뢰할 수 없었습니다.
  • 비유: 마치 범인을 잡으려다 실수로 가장 눈에 띄는 옷을 입은 사람을 지목한 것과 같습니다. AI 가 실제로 그 부분을 보고 판단한 게 아니라, 그냥 "여기가 눈에 띄네"라고 표시한 것뿐입니다.
• 성공한 방법 (Single, LRP, IG):
  • **Single (단일 조각 테스트), LRP(계층별 관련성 전파), IG(적분 기울기)**라는 세 가지 방법이 훨씬 더 AI 의 진짜 생각과 일치했습니다.
  • 비유: 이들은 범인의 지문이나 DNA처럼, AI 가 실제로 판단에 사용한 증거를 정확히 찾아냅니다.

4. 실제 적용 사례: AI 가 발견한 새로운 비밀

이렇게 신뢰할 수 있는 지도를 사용하면 AI 가 숨겨진 비밀을 찾아낼 수 있습니다.

• 사례 1: 유전자 지도 그리기
  • AI 가 조직 슬라이드만 보고 유전자 수치를 예측했는데, 그 '중요한 부분 지도'가 실제 유전자 측정 데이터 (공간 전사체) 와 완벽하게 일치했습니다.
  • 의미: AI 가 단순히 암을 진단하는 것을 넘어, 어떤 조직 부위에서 어떤 유전자가 활발한지를 시각적으로 보여줄 수 있다는 뜻입니다. 마치 조직 안에 유전자의 '지도'를 그려준 것과 같습니다.
• 사례 2: HPV 감염의 숨은 패턴 찾기
  • 목과 머리 부위 암에서 HPV 감염을 예측할 때, AI 는 의사들이 미처 눈치채지 못한 세 가지 다른 패턴으로 감염을 찾아냈습니다.
  • 비유: 의사들은 "염증이 심한 곳"을 보고 진단했지만, AI 는 "염증이 적은 곳"이나 "특정 세포가 모여 있는 곳"을 보고 다른 방식으로 진단했습니다. 이는 AI 가 인간이 알지 못했던 새로운 병리학적 단서를 발견했음을 의미합니다.

5. 결론: 왜 이 연구가 중요한가?

이 논문은 **"AI 의 설명 (히트맵) 을 맹신하지 말고, 검증해야 한다"**는 메시지를 전합니다.

• 기존: "AI 가 빨간색으로 표시했으니 믿자." (위험할 수 있음)
• 새로운 접근: "AI 가 빨간색으로 표시한 게 진짜 이유인지, '조각 지우기' 테스트로 확인하자. 그리고 가장 신뢰할 수 있는 도구 (LRP, Single 등) 를 쓰자."

이 연구를 통해 의료 현장에서 AI 가 내린 판단을 더 신뢰할 수 있게 되었고, AI 가 발견한 새로운 생물학적 단서들을 통해 더 나은 치료법을 개발할 수 있는 길이 열렸습니다. 마치 수사팀이 가짜 단서를 걸러내고 진짜 범인의 흔적을 찾아내는 과정과 같습니다.

기술 요약

이 논문은 다중 인스턴스 학습 (Multiple Instance Learning, MIL) 기반의 디지털 병리학 모델에서 사용되는 **히트맵 (Heatmap)**의 신뢰성과 설명 가능성 (Explainability) 을 평가하기 위한 포괄적인 프레임워크를 제시하고, 이를 통해 기존 방법론들의 한계를 극복하는 새로운 기준을 마련합니다.

주요 내용은 다음과 같습니다.

1. 문제 제기 (Problem)

• 배경: 계산 병리학 (Computational Histopathology) 에서 거대 해상도의 전체 슬라이드 이미지 (WSI) 를 분석할 때, MIL 은 수많은 패치 (패치) 를 집계하여 슬라이드 수준의 예측을 수행하는 핵심 기법입니다.
• 현황: 모델이 어떤 조직 영역에 주목했는지 시각화하기 위해 히트맵이 널리 사용되며, 이는 모델 검증과 생체 표지자 (Biomarker) 발견에 필수적입니다.
• 문제점:
  • 가장 흔하게 사용되는 어텐션 (Attention) 히트맵이 실제로 모델의 의사결정 메커니즘을 얼마나 faithfully(정확하게) 반영하는지에 대한 검증이 거의 이루어지지 않았습니다.
  • 다양한 설명 방법 (Explanation Methods) 이 존재하지만, 어떤 상황에서 어떤 방법을 선택해야 하는지에 대한 체계적인 가이드가 부족합니다.
  • 검증되지 않은 시각화에 의존하면 확인 편향 (Confirmation Bias) 이 발생하거나, 모델이 실제 생물학적 신호가 아닌 아티팩트 (Artifacts) 나 스테이닝 차이에 의존하는 것을 놓칠 수 있습니다.

2. 방법론 (Methodology)

저자들은 추가적인 레이블 없이 MIL 히트맵의 품질을 평가할 수 있는 일반화된 프레임워크를 제안했습니다.

• 패치 플립핑 (Patch Flipping) 기반 평가:
  • 기존 XAI(Explainable AI) 의 'Faithfulness(충실도)' 개념을 MIL 환경에 적용했습니다.
  • 히트맵 점수에 따라 패치를 순서대로 제거 (Perturbation) 하거나 추가하여 모델 예측이 어떻게 변하는지 관찰합니다.
  • 대칭적 관련성 이득 (Symmetric Relevance Gain, SRG) 지수를 계산합니다. 가장 관련성이 높은 패치를 제거할 때 예측이 급격히 변하고, 관련성이 낮은 패치를 제거할 때 예측이 천천히 변할수록 (두 곡선 사이의 면적이 클수록) 히트맵이 더 충실하다고 판단합니다.
• 광범위한 벤치마크:
  • 6 가지 설명 방법 비교: 어텐션 (Attention), 단일 인스턴트 (Single, Perturbation 기반), LRP (Layer-wise Relevance Propagation), IG (Integrated Gradients), G×I, Grad2.
  • 3 가지 MIL 아키텍처: Attention-MIL, Transformer-MIL (TransMIL), Mamba-MIL.
  • 3 가지 작업 유형: 분류 (Classification), 회귀 (Regression), 생존 분석 (Survival).
  • 2 가지 패치 인코더: UNI2, Virchow2 (Foundation Models).
  • 총 10 개의 조직병리학 데이터셋을 대상으로 대규모 실험을 수행했습니다.

3. 주요 기여 (Key Contributions)

1. 검증 프레임워크 제안: 추가 레이블 없이 모델 전략에 대한 히트맵의 충실도를 정량적으로 평가하는 '패치 플립핑' 방법론과 통계적 비교 도구를 개발했습니다.
2. 대규모 벤치마크 결과 도출: 다양한 작업과 아키텍처에 걸쳐 6 가지 설명 방법을 체계적으로 비교하여, 어텐션 히트맵이 대부분의 경우 모델 전략을 반영하지 못함을 입증했습니다.
3. 최적 방법론 가이드라인 제시: 작업 유형과 모델 아키텍처에 따라 가장 적합한 설명 방법이 다르다는 것을 밝혔으며, Single, LRP, IG가 일관되게 우수한 성능을 보임을 확인했습니다.
4. 실제 적용 사례 (Use Cases) 제시:
   • 생물학적 검증: Bulk 유전자 발현 예측 모델의 히트맵을 공간 전사체학 (Spatial Transcriptomics) 데이터와 상관관계 분석하여 생물학적 타당성을 입증했습니다.
   • 모델 전략 발견: 두경부암 (HNSCC) 의 HPV 감염 예측에서 LRP 히트맵을 활용하여, 인간 병리사가 인지하지 못한 새로운 조직학적 패턴 (모델의 다른 예측 전략) 을 발견했습니다.

4. 주요 결과 (Key Results)

• 어텐션 히트맵의 한계: Attention 기반 히트맵은 무작위 베이스라인 (Random Baseline) 과 큰 차이가 없거나, 양/음의 증거를 구분하지 못해 모델의 실제 의사결정 과정을 왜곡하는 경우가 많았습니다.
• 성능 우위 그룹: Single (Perturbation 기반), LRP, IG가 다른 방법들보다 일관되게 높은 충실도 (SRG 점수) 를 보였습니다.
  • Transformer 기반 MIL: LRP 가 가장 우수했습니다.
  • Attention/Mamba 기반 MIL: Single 이 가장 우수했습니다 (특히 회귀 및 생존 분석 작업에서).
• 생물학적 타당성: 충실도가 높은 히트맵 (LRP, Single 등) 은 실제 공간 전사체학 데이터와 높은 상관관계를 보였으며, 이는 모델이 단순한 아티팩트가 아닌 실제 생물학적 신호를 학습했음을 시사합니다.
• 새로운 통찰: HPV 예측 모델이 단순히 세포 크기나 염색체 밀도뿐만 아니라, 종양 내 염증 세포의 분포나 해부학적 위치 (편도 등) 를 기반으로 예측 전략을 다르게 사용한다는 것을 발견했습니다.

5. 의의 및 결론 (Significance)

• 디지털 병리학의 신뢰성 향상: MIL 모델의 설명 가능성을 체계적으로 검증함으로써, 임상 현장에서의 모델 도입 장벽을 낮추고 신뢰도를 높일 수 있습니다.
• 표지자 발견 및 검증: 히트맵을 통해 새로운 생물학적 표지자를 발견하거나, 기존 표지자 예측 모델의 타당성을 공간 전사체학 데이터와 같은 독립적인 지표를 통해 검증하는 새로운 패러다임을 제시했습니다.
• 실용적 가이드: 연구자와 임상가들이 특정 작업 (분류, 회귀, 생존) 과 모델 아키텍처에 맞춰 최적의 설명 방법 (예: Transformer 에는 LRP, Attention/Mamba 에는 Single) 을 선택할 수 있는 구체적인 지침을 제공합니다.

결론적으로, 이 논문은 "단순히 히트맵을 보여주는 것을 넘어, 그 히트맵이 모델의 실제 사고 과정을 얼마나 정확히 반영하는지 검증하는 것"이 디지털 병리학 AI 의 성숙을 위해 필수적임을 강조하며, 이를 위한 표준화된 평가 체계와 최적의 방법론을 제시합니다.