[논문 리뷰] Evolution-based Fine-tuning of CNNs for Prostate Cancer Detection

Khashayar Namdar, Isha Gujrathi외 2명|arXiv (Cornell University)|2019. 11. 04.

Advanced Neural Network Applications참고 문헌 20인용 수 4

한 줄 요약

이 논문은 6채널의 확산 영상 자기공명영상(DW-MRI)에서 전립선암 검출을 위해 사전 훈련된 컨볼루션 네트워크(CNN)의 완전히 연결된(FC) 계층을 최적화하기 위해 유전 알고리즘(GA)을 사용하는 진화 기반 미세조정 방법을 제안한다. 개별 가중치가 아닌 계층 수준에서 GA를 적용함으로써, 테스트 AUC를 1,334개의 테스트 슬라이스에서 0.707에서 0.773으로 9.3% 향상시켰으며, 특성 추출 계층을 재학습하지 않고도 뚜렷한 성능 향상을 달성하였다.

ABSTRACT

Convolutional Neural Networks (CNNs) have been used for automated detection of prostate cancer where Area Under Receiver Operating Characteristic (ROC) curve (AUC) is usually used as the performance metric. Given that AUC is not differentiable, common practice is to train the CNN using a loss functions based on other performance metrics such as cross entropy and monitoring AUC to select the best model. In this work, we propose to fine-tune a trained CNN for prostate cancer detection using a Genetic Algorithm to achieve a higher AUC. Our dataset contained 6-channel Diffusion-Weighted MRI slices of prostate. On a cohort of 2,955 training, 1,417 validation, and 1,334 test slices, we reached test AUC of 0.773; a 9.3% improvement compared to the base CNN model.

연구 동기 및 목표

• 표준 역전파 알고리즘과 미분 가능한 손실 함수를 사용함에도 불구하고 CNN 기반 전립선암 검출에서 AUC 성능이 최적화되지 않는 문제를 해결하기 위해.
• SGD가 AUC와 같은 비미분 가능한 지표를 최적화하는 데 한계가 있음을 고려하여, 분류기 헤드에 진화 계산을 도입하기 위해.
• 개별 가중치가 아닌 계층 수준에서 작동함으로써 GA 기반 CNN 최적화의 계산 비용을 줄이기 위해.
• 완전히 연결된 계층의 진화 기반 미세조정이 특성 추출 기반을 동결한 상태에서도 AUC를 뚜렷이 향상시킬 수 있음을 입증하기 위해.
• 실제 DW-MRI 데이터셋(414명 환자에서의 5,706개 전립선 슬라이스)을 기반으로 방법을 검증하기 위해.

제안 방법

• 3개의 컨볼루션 계층과 3개의 완전히 연결된(FC) 계층을 갖는 사전 훈련된 CNN을 사용하였으며, Xavier 및 균일 분포를 통해 가중치를 초기화하였다.
• CNN은 고정된 학습률, 모멘터리, L2 정규화를 사용한 교차 엔트로피 손실과 SGD를 통해 먼저 훈련되었다.
• 유전 알고리즘이 FC 계층에만 적용되었으며, 각 계층을 단일 개체로 간주하여 교배 및 돌연변이 연산을 수행하였다.
• 초기 GA 집단은 512개의 분류기로 구성되었으며, 이 중 하나는 훈련된 CNN에서 유래하였고, 나머지 511개는 무작위로 초기화되었다.
• 적합도 평가가 훈련 세트의 AUC 기반으로 이루어졌으며, 세대에 걸쳐 선택, 교배, 돌연변이(1% 확률)가 적용되었다.
• 검증 및 훈련 AUC가 더 이상 향상되지 않을 경우 최적화를 중단하였으며, 최종 모델는 검증 AUC가 가장 높은 모델로 선택되었다.

실험 결과

연구 질문

• RQ1 유전 알고리즘이 DW-MRI에서 전립선암 검출의 AUC를 향상시키기 위해 CNN의 완전히 연결된 계층을 효과적으로 최적화할 수 있는가?
• RQ2 노드 수준 최적화와 비교해 계층 수준의 GA 돌연변이 및 교배가 성능 향상을 유지하면서 계산 비용을 줄일 수 있는가?
• RQ3 특성 추출 계층이 동결된 상태에서 진화 기반 미세조정이 표준 SGD 훈련보다 AUC에서 뛰어난 성능을 낼 수 있는가?
• RQ4 AUC 향상은 훈련, 검증 및 독립 테스트 세트 전반에 걸쳐 통계적으로 유의미하고 견고한가?
• RQ5 최종 모델 성능 측면에서 기존의 AUC 모니터링을 통한 전통적 훈련 방식과 비교해 본 논문의 방법은 어떠한가?

주요 결과

• 제안된 GA 기반 미세조정으로 테스트 AUC가 SGD로 훈련된 모델의 0.707에서 0.773으로 9.3% 상승하였다.
• 최고의 모델은 GA의 단지 3세대 만에 도달하였으며, 검증 AUC는 0.815에 도달했고, 테스트 AUC는 0.773였다.
• GeForce GTX 1060 GPU에서 단 10분 만에 최적 성능을 달성하여 계산 가능성의 가능성을 입증하였다.
• GA는 훈련 세트의 AUC를 SGD의 0.867에서 0.877로 향상시켜 더 나은 일반화 성능을 나타내었다.
• 역전파 제약 조건을 우회하기 위해 진화적 탐색을 통해 비미분 가능한 AUC 지표를 성공적으로 최적화하였다.
• 계층 수준의 GA 접근 방식은 노드 수준 최적화 대비 계산 오버헤드를 줄여 저사양 하드웨어에서도 실용적으로 적용 가능하게 하였다.