[논문 리뷰] Automatic Data Augmentation via Deep Reinforcement Learning for Effective Kidney Tumor Segmentation

Tiexin Qin, Ziyuan Wang외 4명|arXiv (Cornell University)|2020. 02. 22.

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

한 줄 요약

이 논문은 CT 영상에서 신장 종양 세분화를 위한 딥 강화학습(DRL) 기반의 자동 데이터 증강 방법을 제안한다. 여기서 에이전트는 검증 세트에서 DSC 향상도를 최대화함으로써 종료형으로 최적의 증강 정책을 학습한다. 이 방법은 기존의 방법과 GAN 기반 증강 기법을 모두 능가하며, mIoU 75.8%와 DSC 84.0%를 달성하여, 제한된 데이터에서 뛰어난 경계 국소화 능력과 강건성을 보여준다.

ABSTRACT

Conventional data augmentation realized by performing simple pre-processing operations (\eg, rotation, crop, \etc) has been validated for its advantage in enhancing the performance for medical image segmentation. However, the data generated by these conventional augmentation methods are random and sometimes harmful to the subsequent segmentation. In this paper, we developed a novel automatic learning-based data augmentation method for medical image segmentation which models the augmentation task as a trial-and-error procedure using deep reinforcement learning (DRL). In our method, we innovatively combine the data augmentation module and the subsequent segmentation module in an end-to-end training manner with a consistent loss. Specifically, the best sequential combination of different basic operations is automatically learned by directly maximizing the performance improvement (\ie, Dice ratio) on the available validation set. We extensively evaluated our method on CT kidney tumor segmentation which validated the promising results of our method.

연구 동기 및 목표

• 제한된 레이블이 부여된 훈련 데이터로 인해 의료 영상 세분화 성능이 저하되는 문제를 해결하기 위해.
• 고정된 랜덤 연산에 의존하고 세분화 성능 피드백이 없는 기존 데이터 증강의 한계를 극복하기 위해.
• 강화학습을 활용해 이미지별로 최적의 증강 정책을 탐색함으로써, 데이터 증강과 세분화를 동시에 최적화하는 종료형 학습 프레임워크를 개발하기 위해.
• 증강 정책 학습 과정에서 세분화 성능(Dice 점수)을 직접 최대화하여, 생성된 데이터가 후행 작업 정확도를 향상시킴을 보장하기 위해.
• 실제 임상 CT 신장 종양 세분화 데이터셋을 대상으로 본 방법을 검증하고, 최신 기준 증강 기반 방법들과 비교하기 위해.

제안 방법

• 본 방법은 데이터 증강을 마르코프 결정 과정(MDP)으로 모델링하며, 에이전트는 세분화 성능 향상을 위해 이미지 변환(예: 반전, 자르기, 회전)을 선택한다.
• 최적의 증강 작업 순서를 학습하기 위해 듀얼링 딥 Q-네트워크(Dueling DQN)를 사용하며, 보상 함수는 검증 세트에서 DSC 향상도로 정의된다.
• 세분화 모델을 대체 정책으로 사용하여 DRL 훈련 중 피드백을 제공함으로써, 증강과 세분화를 종료형으로 공동 최적화할 수 있도록 한다.
• 에이전트는 환경(CT 영상 및 마스크)과 상호작용하며 입력 이미지를 변형하는 동작을 선택하고, 결과 세분화 품질에 따라 보상을 수신한다.
• 최종 증강 데이터셋은 원본 이미지와 DRL에 의해 생성된 증강 이미지를 혼합하여 구성되며, 훈련 세트 크기를 두 배로 늘린다.
• 모든 방법 간 공정한 비교를 위해, Adam 옵timizer를 사용하여 공통된 초모수를 공유해 프레임워크를 훈련시킨다.

실험 결과

연구 질문

• RQ1 강화학습 에이전트는 제한된 의료 영상 데이터에서 세분화 성능을 최대화하는 이미지별로 특화된 증강 순서를 학습할 수 있는가?
• RQ2 세분화와 증강을 종료형으로 공동 훈련하는 것이 증강 모듈과 세분화 모듈를 별도로 훈련하는 것보다 성능을 향상시키는가?
• RQ3 제안된 DRL 기반 증강 방법은 신장 종양 세분화에서 기존의 증강 기법과 GAN 기반 데이터 생성 기법과 비교해 어떻게 성능을 내는가?
• RQ4 에이전트가 자주 선택하는 증강 작업 유형(예: 반전, 자르기)은 무엇이며, 이들은 종양 경계의 세분화 향상과 관련이 있는가?
• RQ5 이 방법은 레이블이 제한된 실세계 임상 데이터셋에 일반화되어 최신 기술 수준의 성능을 달성할 수 있는가?

주요 결과

• 제안된 방법은 mIoU 75.8%를 달성하여, 기존 증강 기법(74.3%)과 다른 기준 기반 방법들을 크게 앞서며 유의미한 성능 향상을 보였다.
• DSC는 84.0%를 기록하여, 다음으로 우수한 성능을 보인 방법(기존 증강 기반 83.2%)을 능가했으며, 민감도(94.0%)와 특이도(PPV: 80.2%)도 가장 높았다.
• 중심점 거리(CD)는 8.86 mm로 감소하였고, 하우스도르프 거리(HD)는 8.52 mm, 평균 표면 거리(ASD)는 4.01 mm로 감소하여, 뛰어난 경계 국소화 정확도를 보였다.
• 에이전트는 수평 반전과 자르기 작업을 자주 선택하였으며, 이는 종양의 가시성을 향상시키고, 크거나 비정상적인 형태의 종양에 대한 일반화 능력을 향상시켰다.
• 시각적 결과는 본 방법이 특히 대trast가 낮거나 비정상적인 형태를 가진 영역에서 더 정확한 종양 경계 예측을 생성함을 확인했다.
• 제거 실험 결과, 세분화 모듈의 피드백을 통한 종료형 공동 학습이 고립된 증강 또는 고정된 작업 집합보다 더 나은 정책 학습을 이끌었다.