[논문 리뷰] Partial transfusion: on the expressive influence of trainable batch norm parameters for transfer learning

Fahdi Kanavati, Masayuki Tsuneki|arXiv (Cornell University)|2021. 02. 10.

COVID-19 diagnosis using AI참고 문헌 25인용 수 14

한 줄 요약

이 논문은 의료 영상 분야의 전이 학습에서 훈련 가능한 배치 정규화(BN) 파라미터—스케일 γ와 시프트 β—만을 미세조정하는 방법을 제안하며, 이 방법이 전체 모델 미세조정과 비교해 유사한 성능을 달성함을 입증한다. 이 방법은 더 적은 파라미터와 더 높은 학습률 덕분에 더 빠른 수렴을 가능하게 하며, 네 가지 아키텍처와 아홉 개인 의료 영상 데이터셋에서 일관된 성능을 보였다.

ABSTRACT

Transfer learning from ImageNet is the go-to approach when applying deep learning to medical images. The approach is either to fine-tune a pre-trained model or use it as a feature extractor. Most modern architecture contain batch normalisation layers, and fine-tuning a model with such layers requires taking a few precautions as they consist of trainable and non-trainable weights and have two operating modes: training and inference. Attention is primarily given to the non-trainable weights used during inference, as they are the primary source of unexpected behaviour or degradation in performance during transfer learning. It is typically recommended to fine-tune the model with the batch normalisation layers kept in inference mode during both training and inference. In this paper, we pay closer attention instead to the trainable weights of the batch normalisation layers, and we explore their expressive influence in the context of transfer learning. We find that only fine-tuning the trainable weights (scale and centre) of the batch normalisation layers leads to similar performance as to fine-tuning all of the weights, with the added benefit of faster convergence. We demonstrate this on a variety of seven publicly available medical imaging datasets, using four different model architectures.

연구 동기 및 목표

• 의료 영상 전이 학습에서 훈련 가능한 배치 정규화(BN) 파라미터의 표현적 영향을 조사한다.
• 최적의 성능을 위해 전체 모델 미세조정이 필수적이라는 일반적인 가정을 도전한다.
• 전체 미세조정과 동등한 성능을 낼 수 있는지 확인하기 위해 BN 애페인 파라미터(γ 및 β)만 업데이트하는 것이 가능한지 탐색한다.
• 전이 학습 중 BN 파라미터 업데이트와 고정된 이동 평균 간의 영향을 평가한다.
• BN 레이어의 내재적 표현력 이해를 위해 랜덤 가중치를 사용하면서도 오직 BN 파라미터만 훈련시키는 것이 가능한지 평가한다.

제안 방법

• 모든 다른 가중치를 고정한 채로 배치 정규화 레이어의 훈련 가능한 스케일(γ) 및 시프트(β) 파라미터만 미세조정한다.
• 다섯 가지 훈련 전략—특징 추출, 전체 미세조정, BN 전용 미세조정, 이동 평균 업데이트(BMA), 랜덤 가중치에 BN 전용 훈련—간의 성능을 비교한다.
• 아홉 개인 의료 영상 데이터셋에서 네 가지 사전 훈련된 CNN 아키텍처—DenseNet121, ResNet50V2, InceptionV3, EfficientNetB3—를 사용한다.
• 최종 완전 연결 레이어를 먼저 훈련한 후 BN 파라미터를 미세조정하거나, 둘 다 동시에 훈련한다.
• 훈련 중에 BN 레이어를 추론 모드로 유지하여 이동 평균 업데이트로 인한 불안정성을 방지한다.
• 다중 레이블 및 이진 분류 작업에서 테스트 세트의 ROC AUC를 사용해 모델을 평가한다.

실험 결과

연구 질문

• RQ1 의료 영상 분류에서 배치 정규화 애페인 파라미터(γ 및 β)만 미세조정하는 것이 전체 모델 미세조정과 비교해 유사한 성능을 달성할 수 있는가?
• RQ2 미세조정 중에 배치 정규화 이동 평균을 업데이트하는 것이 이동 평균을 고정한 경우에 비해 성능을 떨어뜨리는가?
• RQ3 랜덤 초기화된 가중치를 사용하면서도 오직 BN 파라미터만 훈련시키는 경우 성능는 어떻게 되는가?
• RQ4 BN 전용 미세조정이 수렴 속도와 모델 안정성에 어떤 영향을 미치는가?
• RQ5 BN 전용 미세조정의 효과성에 영향을 미치는 아키텍처나 데이터셋 특성 요소가 존재하는가?

주요 결과

• 모든 네 개 아키텍처와 아홉 개 데이터셋에서 배치 정규화 애페인 파라미터(γ 및 β)만 미세조정하는 것이 전체 모델 미세조정과 유사한 성능을 달성한다.
• DenseNet121에서 BN 전용 미세조정은 모든 데이터셋에서 전체 미세조정 대비 ROC AUC 점수를 1~2% 이내로 유지했으며, 어떤 작업에서도 유의미한 성능 저하가 없었다.
• BN 전용 미세조정은 더 적은 훈련 가능한 파라미터와 더 높은 학습률 사용 덕분에 더 빠른 수렴을 이끌었다.
• 전체 모델을 특징 추출기로 사용하는 경우(미세조정 없이)는 항상 전체 및 BN 전용 미세조정보다 성능이 열 劣했다.
• 이동 평균 통계치를 업데이트하는 것(BMA)은 대부분의 데이터셋에서 성능 저하를 초래했으며, 이는 이동 평균을 고정하는 것이 유리함을 확인시켰다.
• 랜덤 가중치에 BN 파라미터만 훈련시킨 경우, 흉부 X-ray나 말라리아와 같은 간단한 데이터셋에서는 만족스러운 성능를 달성했지만, EfficientNetB3에서는 기울기 포화로 인해 훈련이 불가능했으며, 이는 아키텍처 의존성을 강조했다.