DiNTS: 3차원 의료 영상 분할을 위한 미분 가능한 신경망 구조 탐색

최근 들어 신경망 아키텍처 탐색(NAS) 기술이 의료 영상 분할을 위한 고성능 네트워크를 자동으로 탐색하는 데 적용되고 있다. 일반적으로 NAS의 탐색 공간은 셀 간의 연결을 제어하는 네트워크 아키텍처 수준과 각 셀 내부의 연산을 결정하는 셀 수준으로 구성된다. 기존의 방법들은 대규모 3차원(3D) 영상 데이터셋을 대상으로 할 경우 긴 탐색 시간이 소요되거나, 사전 정의된 아키텍처(예: U형 또는 단일 경로)에 한정되는 한계를 가지고 있다. 본 연구에서는 3D 의료 영상 분할에서 NAS의 세 가지 중요한 측면—유연한 다중 경로 네트워크 아키텍처, 높은 탐색 효율성, 제한된 GPU 메모리 사용—에 초점을 맞춘다. 우리는 매우 유연한 네트워크 아키텍처 탐색 공간 내에서 빠른 기울기 기반 탐색을 가능하게 하는 새로운 미분 가능(differentiable) 탐색 프레임워크를 제안한다. 그러나 미분 가능 방식에서 탐색된 최적의 연속 모델을 이산화할 경우, 최종 이산 모델이 하위 최적(서브옵티멀)이 되는 '이산화 갭(discretization gap)' 문제가 발생할 수 있다. 이를 해결하기 위해 우리는 아키텍처 손실( topology loss)을 제안한다. 또한, 탐색 과정 중에 탐색된 3D 모델의 GPU 메모리 사용량을 예산 제약 조건 내에서 제한한다. 제안한 미분 가능 네트워크 아키텍처 탐색 기법(DiNTS, Differentiable Network Topology Search)은 10개의 도전적인 분할 과제를 포함하는 Medical Segmentation Decathlon(MSD) 챌린지에서 평가되었으며, 최신 기술 수준(SOTA)의 성능을 달성하고 MSD 챌린지 랭킹에서 최상위 순위를 기록하였다.