[논문 리뷰] iUNets: Fully invertible U-Nets with Learnable Up- and Downsampling

Christian Etmann, Rihuan Ke외 1명|arXiv (Cornell University)|2020. 05. 11.

Parallel Computing and Optimization Techniques참고 문헌 25인용 수 13

한 줄 요약

이 논문은 지수화된 반대칭 행렬 기반의 학습 가능한 역행성 상향 및 하향 샘플링 연산을 갖춘 완전히 역행성인 U-Net 아키텍처인 iUNets를 소개한다. 역행성 덕분에 메모리 효율적인 역전파가 가능해져, GPU 메모리 제약 하에서 더 깊고 더 큰 네트워크를 훈련시킬 수 있으며, 3D 의료 영상 분할 및 CT 후처리에서 표준 U-Net을 능가한다. 또한 경쟁적인 우도 성능을 보이며 정규화 흐름으로도 활용 가능하다.

ABSTRACT

U-Nets have been established as a standard architecture for image-to-image learning problems such as segmentation and inverse problems in imaging. For large-scale data, as it for example appears in 3D medical imaging, the U-Net however has prohibitive memory requirements. Here, we present a new fully-invertible U-Net-based architecture called the iUNet, which employs novel learnable and invertible up- and downsampling operations, thereby making the use of memory-efficient backpropagation possible. This allows us to train deeper and larger networks in practice, under the same GPU memory restrictions. Due to its invertibility, the iUNet can furthermore be used for constructing normalizing flows.

연구 동기 및 목표

• 대규모 3D 의료 영상에서 표준 U-Net의 높은 메모리 소비 문제를 해결하기 위해 메모리 효율적인 훈련을 가능하게 하는 것.
• 중간 활성화 값을 저장하지 않고도 역전파를 지원하는 완전히 역행성인 U-Net 아키텍처를 개발하는 것.
• 고정된 GPU 메모리 제약 조건 하에서 더 크고 깊은 네트워크를 허용하는 것.
• 이 아키텍처가 생성 모델링을 위한 정규화 흐름으로 사용 가능한지 확인하는 것.
• 체적 분할 및 CT 재구성과 같은 3D 이미지 간 작업에서 성능 향상을 이루는 것.

제안 방법

• 지수화된 반대칭 컨벌루션 커널을 사용한 학습 가능한 역행성 하향 샘플링을 제안하며, 이는 전단사적이고 미분 가능하다.
• 학습된 하향 샘플링의 역연산으로서의 역행성 상향 샘플링을 정의하여 완전한 역행성 보장.
• 반대칭 행렬에서 유도된 직교 컨벌루션 연산자를 사용하여 훈련 중 역행성과 안정성을 유지.
• 공간적 하향 샘플링을 채널 확장으로 매핑하기 위해 구조화된 커널 재정렬 메커니즘을 도입하여 정보 손실 방지.
• 활성화 재구성에 의한 메모리 효율적인 역전파를 사용하여 3D 의료 영상 작업에 아키텍처 적용.
• 아핀 커플링 레이어를 사용해 iUNet 블록을 스택하여 정규화 흐름을 구성함으로써 우도 기반 생성 모델링 가능.

실험 결과

연구 질문

• RQ1 U-Net 아키텍처에서 최대풀링 및 역전치 컨벌루션과 같은 비역행성 연산을 대체할 수 있는 학습 가능한 역행성 상향/하향 샘플링 연산을 설계할 수 있는가?
• RQ2 결과적으로 도출된 완전히 역행성인 U-Net(iUNet)이 고정된 GPU 메모리 조건 하에서 더 깊고 더 큰 네트워크의 메모리 효율적인 훈련을 가능하게 하는가?
• RQ3 iUNets가 표준 U-Net에 비해 3D 의료 영상 분할 및 후처리에서 뛰어난 성능을 내는가?
• RQ4 iUNet 아키텍처는 정규화 흐름으로 사용 가능한가? 그리고 최첨단 방법과 비교해 우도 성능은 어떠한가?
• RQ5 학습된 역행성 연산이 아티팩트를 유발하지 않고 다양한 이미지 간 작업에 일반화 가능한가?

주요 결과

• 64채널을 갖는 iUNet는 BraTS2018 검증 세트에서 0.850의 Dice 스코어를 기록하여 표준 U-Net(0.833)과 더 큰 iUNets(예: iUNet-64: 0.850 vs. U-Net: 0.833)를 능가했다.
• 3D CT 후처리 작업에서 iUNet는 베이스라인 U-Net보다 뚜렷이 뛰어난 성능을 보이며, 메모리 효율성 덕분에 더 큰 모델 용량을 활용할 수 있음을 입증했다.
• iUNet를 사용한 메모리 효율적인 역전파가 표준 역전파와 유사한 손실 곡선을 생성하여 훈련 안정성을 확인했다.
• iUNet 기반 정규화 흐름은 CIFAR10에서 3.60 비트/차원의 음의 로그우도를 기록했으며, Real NVP(3.49 비트/차원)에 비해 약간 열등하여 향상 여지가 있음을 시사했다.
• 역행성 확보를 위해 유한 급수 형태의 행렬 지수를 사용함으로써 우도 추정에 잠재적 편향이 발생했으며, 이는 성능 격차의 원인일 수 있다.
• iUNet 아키텍처는 메모리 제약 조건 하에서도 더 큰 네트워크 훈련이 가능하며, 이는 채널 수가 증가함에 따라 성능 향상(예: iUNet-64)을 통해 입증되었다.