[논문 리뷰] GradMask: Reduce Overfitting by Regularizing Saliency

Becks Simpson, Francis Dutil외 2명|arXiv (Cornell University)|2019. 04. 16.

Radiomics and Machine Learning in Medical Imaging참고 문헌 11인용 수 14

한 줄 요약

GradMask는 경량 의료 영상 데이터셋에서 과적합을 줄이기 위해 기울기 기반 기여도를 사용하여 병변 영역 외부의 특징에 중요도를 할당하는 시각화 지도를 정규화하는 새로운 정규화 방법이다. 이는 모델의 주의 집중이 진단 기준 세그멘테이션과 일관성을 유지하도록 유도함으로써 테스트 AUC를 1–4% 향상시킨다.

ABSTRACT

With too few samples or too many model parameters, overfitting can inhibit the ability to generalise predictions to new data. Within medical imaging, this can occur when features are incorrectly assigned importance such as distinct hospital specific artifacts, leading to poor performance on a new dataset from a different institution without those features, which is undesirable. Most regularization methods do not explicitly penalize the incorrect association of these features to the target class and hence fail to address this issue. We propose a regularization method, GradMask, which penalizes saliency maps inferred from the classifier gradients when they are not consistent with the lesion segmentation. This prevents non-tumor related features to contribute to the classification of unhealthy samples. We demonstrate that this method can improve test accuracy between 1-3% compared to the baseline without GradMask, showing that it has an impact on reducing overfitting.

연구 동기 및 목표

• 비생물학적 특징(예: 스캐너 아티팩트)과의 허위 상관관계를 학습하는 소규모 의료 영상 데이터셋에서의 과적합 문제를 해결한다.
• 분류 과정에서 병변 경계 외부의 특징에 의존하는 것을 명시적으로 막음으로써 모델의 일반화 능력을 향상시킨다.
• 기울기 기반 시각화 지도를 활용하여 잘못된 특징 기여도를 페널티 부여하는 정규화 방법을 개발한다.
• 진단 기준 세그멘테이션과 일치하는 시각화 지도의 정규화가 테스트 성능 향상에 기여함을 입증한다.

제안 방법

• 비병변 영역에 대해 이진 세그멘테이션 마스크를 사용해 마스킹된 입력 픽셀에 대한 예측 클래스 점수의 기울기 L2 노름을 페널티로 부여하는 손실 항목을 도입한다.
• 비건강 상태 예측($\hat{y}_1$)의 입력 $\mathbf{x}$에 대한 기울기를 시각화 지도로 사용하고, 병변 외부의 기울기를 억제하기 위해 마스크 $1 - \mathbf{x}_{\text{seg}}$를 적용한다.
• 총 손실를 $\mathcal{L} = \sum_{\mathbf{x} \in D} \mathcal{L}_c + \left\| \frac{\partial \hat{y}_1}{\partial \mathbf{x}} \cdot (1 - \mathbf{x}_{\text{seg}}) \right\|_2^2$ 로 정의하며, 여기서 $\mathcal{L}_c$ 는 표준 분류 손실이다.
• 건강 및 비건강 예측 간 절대 차이의 기울기를 계산하여 특징을 구분하는 데 중점을 두는 '대비' 시각화 지도 버전을 탐색한다.
• 의료 영상 벤치마크(MSD 및 BraTS)에서 미세조정된 소규모 데이터셋을 사용해 학습 중에 이 방법을 적용하고, 검증 AUC를 활용해 하이퍼파rameter를 튜닝한다.
• 입력 기울기 시각화 지도를 사용해 병변과 관련 없는 특징을 식별하고 정규화함으로써, 주의 집중이 생물학적으로 관련된 영역에 국한되도록 보장한다.

실험 결과

연구 질문

• RQ1 소규모 의료 영상 데이터셋에서 비생물학적 특징에 과적합하는 경향이 높을 때, 시각화 지도의 명시적 정규화가 일반화 성능 향상에 기여하는가?
• RQ2 병변 경계 외부의 기울기 기반 특징 중요도를 페널티로 줄일 경우, 기준 모델 대비 테스트 AUC에서 측정 가능한 향상이 발생하는가?
• RQ3 학습 데이터가 제한적일 경우(예: 64–128개 샘플), GradMask가 스캐너 특화 아티팩트를 기억하는 것과 같은 과적합을 얼마나 효과적으로 줄이는가?
• RQ4 건강 및 비건강 예측 간 차이를 강조하는 대비 기반 시각화 지도 버전이 클래스별로 정규화된 기존 방법보다 더 높은 성능을 내는가?
• RQ5 이전 연구(Zech et al., 2018)에서 보고된 바와 같이, 병원 특화 영상 아티팩트와 같은 허위 상관관계에 대한 의존도가 GradMask로 얼마나 줄어드는가?

주요 결과

• GradMask는 간, 췌장, 뇌, 심장 등 네 가지 의료 영상 데이터셋에서 정규화 없이 학습한 기준 모델 대비 테스트 AUC를 1–4% 향상시켰다.
• 낮은 샘플 수에서도 측정 가능한 향상이 관찰되었으며, 첫 번째 양성 결과는 64–128개의 훈련 샘플에서 관찰되었다.
• BraTS 뇌 종양 데이터셋에서 GradMask는 기준 모델의 테스트 AUC 0.826 ± 0.026에서 0.798 ± 0.019로 감소했지만, 다른 데이터셋에서는 일관되게 향상이 관찰되었다.
• 대비 기반 시각화 지도 버전은 여러 시드와 데이터셋에서 더 높은 테스트 AUC를 기록함으로써 일반화 능력에서 표준 버전을 능가했다.
• 검증 AUC를 활용한 하이퍼파rameter 튜닝을 통해 최적의 설정을 도출했고, MSD 세 개 및 BraTS 한 개 데이터셋에서 20개의 랜덤 시드 전반에 걸쳐 일관되게 성능 향상을 보였다.
• 모델이 훈련 데이터에서 클래스 레이블과 통계적으로 상관관계가 있는 병변 외부 영역의 특징에 중요도를 할당하는 것을 억제함으로써 과적합을 효과적으로 줄였다.