[논문 리뷰] Distributed Weight Consolidation: A Brain Segmentation Case Study

Patrick McClure, Charles Zheng외 6명|PubMed|2018. 05. 28.

Domain Adaptation and Few-Shot Learning참고 문헌 30인용 수 5

한 줄 요약

이 논문은 지속적 학습 방법인 분산 가중치 통합(DWC)을 제안한다. 이 방법은 개인정보, 윤리적 또는 법적 제약으로 인해 공유되지 않는 독립적인 비중叠된 sMRI 데이터셋에서 사전 학습된 깊이 신경망을 통합하여, 순차적 학습이나 모델 앙상블 없이 뇌 분할을 수행한다. DWC는 개별 기관의 데이터에서 일반화 능력을 향상시키며, 대규모 독립적인 다기관 데이터셋에서 뛰어난 성능을 유지한다. 또한 순차적 학습 기반의 앙상블 기준 모델보다 기관별 테스트 성능에서 뛰어나며 일반화 능력을 유지한다.

ABSTRACT

Collecting the large datasets needed to train deep neural networks can be very difficult, particularly for the many applications for which sharing and pooling data is complicated by practical, ethical, or legal concerns. However, it may be the case that derivative datasets or predictive models developed within individual sites can be shared and combined with fewer restrictions. Training on distributed data and combining the resulting networks is often viewed as continual learning, but these methods require networks to be trained sequentially. In this paper, we introduce distributed weight consolidation (DWC), a continual learning method to consolidate the weights of separate neural networks, each trained on an independent dataset. We evaluated DWC with a brain segmentation case study, where we consolidated dilated convolutional neural networks trained on independent structural magnetic resonance imaging (sMRI) datasets from different sites. We found that DWC led to increased performance on test sets from the different sites, while maintaining generalization performance for a very large and completely independent multi-site dataset, compared to an ensemble baseline.

연구 동기 및 목표

• 개인정보, 윤리적 또는 법적 제약으로 인해 공유되지 않는 분산된 임상 뇌 영상 데이터에서 깊이 신경망을 학습하는 데 도전하는 것.
• 중앙집중식 데이터 접근 없이도 순차적 학습이 필요 없이 다수 기관에서 독립적으로 학습된 모델의 지식을 통합하는 방법을 개발하는 것.
• 개별 기관의 테스트 세트에서 모델의 일반화 능력을 향상시키면서도, 완전히 독립적인 다기관 대규모 데이터셋에서의 성능를 유지하는 것.
• 모델 앙상블과 순차적 지속적 학습의 확장 가능한 대안을 제공함으로써 분산된 의료 영상 응용 분야에 기여하는 것.

제안 방법

• DWC는 다양한 sMRI 데이터셋에서 독립적으로 학습된 3D 확장 컨볼루션 신경망(MeshNet 기반)의 가중치를 통합하기 위해 변동성 추론 프레임워크를 사용한다.
• 이를 위해 네트워크 가중치에 대한 사후 분포로 통합을 공식화하고, 효율적인 최적화를 위해 평균장 근사법을 적용한다.
• 모델 간 지식 이동을 가능하게 하기 위해 공통된 사전 분포를 가중치에 적용함으로써 직접적인 데이터 공유 없이도 지식 공유가 가능하다.
• 추론 시 기대 출력을 근사하기 위해 몬테카를로 샘플링(10개 샘플)을 적용한 확률적 경량 최적화 방법을 사용한다.
• 기존 지속적 학습과 달리 순차적이지 않은 통합이 가능하여, 여러 기관의 모델을 어떤 순서로든 조합할 수 있다.
• 최종 추론 네트워크로 통합된 모델의 최대사후확률( MAP ) 추정치를 사용하며, 이를 보류된 테스트 세트에서 평가한다.

실험 결과

연구 질문

• RQ1 비중엽된 sMRI 데이터셋에서 독립적으로 학습된 깊이 신경망의 지식을 효과적으로 통합하여 개별 기관의 테스트 세트 성능을 향상시킬 수 있는가?
• RQ2 기존 기준 모델 대비, 통합된 모델이 완전히 독립적인 다기관 대규모 데이터셋에서 일반화 성능를 유지하는가?
• RQ3 모델 앙상블과 순차적 지속적 학습 대비 DWC의 성능 및 데이터셋 순서에 대한 강건성은 어떠한가?
• RQ4 데이터 공유가 제한되지만 모델 공유가 가능한 의료 영상 작업에 대해 DWC가 효과적으로 적용될 수 있는가?
• RQ5 비순차적 통합이 다양한 뇌 영상 기관에서의 모델 성능에 어떤 영향을 미치는가?

주요 결과

• DWC는 HCP, NKI, Buckner, WU120 데이터셋에서 가중 평균 DICE 스코어 78.30을 기록하여 앙상블 기준 모델보다 유의미하게 높은 성능( p = 1.66e-15 )을 보였다.
• ABIDE 데이터셋에서 DWC의 DICE 스코어(70.76)는 앙상블과 유의미한 차이가 없었음(p = 0.733)을 보여, 독립된 데이터에서 일반화 능력이 떨어지지 않았음을 시사한다.
• DWC는 VCL와 유사한 성능를 보였지만, 순차적 지속적 학습 대비 데이터셋 순서에 덜 민감했으며, 더 큰 데이터셋을 마지막에 학습했을 경우 78.28의 가중 평균 스코어를 기록하여 순차적 학습의 75.95보다 높았다.
• 시각적 점검 결과, DWC의 분할 결과는 참값과 유사했으며, 주로 영역 경계에서 오류가 발생했고, 이는 고성능 HNBW_MAP 모델과 유사했다.
• 이 방법은 HCP, NKI, Buckner, WU120의 네 개 기관에서 유도된 모델을 하나의 모델로 통합하여 모든 기관과 독립된 데이터셋에서 잘 일반화되는 것을 성공적으로 달성했다.
• 몬테카를로 샘플링을 통한 변동성 추론은 통합된 모델의 기대 출력을 안정적으로 추정함으로써 신뢰할 수 있는 추론을 가능하게 했다.