[논문 리뷰] Symmetric block-low-rank layers for fully reversible multilevel neural networks

Bas Peters, Eldad Haber외 1명|arXiv (Cornell University)|2019. 12. 14.

Sparse and Compressive Sensing Techniques참고 문헌 21인용 수 4

한 줄 요약

이 논문은 영상과 같은 대규모 3D 데이터를 위한 완전히 가역적인 다수준 신경망을 가능하게 하기 위해 대칭 블록-저랭크(BLR) 레이어를 제안한다. 이는 가역적 다운샘플링 과정에서 채널 수가 증가함에 따라 발생하는 메모리 폭발 문제를 해결한다. 컨볼루션 커널을 대칭 BLR 구조로 인수분해함으로써, 파라미터 저장 용량을 극적으로 줄이면서도 가역성과 전체 영상 입력에 대한 엔드 투 엔드 학습을 유지한다. 이는 기존의 표준 가역 신경망에서는 이전에 불가능했던 일이다.

ABSTRACT

Factors that limit the size of the input and output of a neural network include memory requirements for the network states/activations to compute gradients, as well as memory for the convolutional kernels or other weights. The memory restriction is especially limiting for applications where we want to learn how to map volumetric data to the desired output, such as video-to-video. Recently developed fully reversible neural networks enable gradient computations using storage of the network states for a couple of layers only. While this saves a tremendous amount of memory, it is the convolutional kernels that take up most memory if fully reversible networks contain multiple invertible pooling/coarsening layers. Invertible coarsening operators such as the orthogonal wavelet transform cause the number of channels to grow explosively. We address this issue by combining fully reversible networks with layers that contain the convolutional kernels in a compressed form directly. Specifically, we introduce a layer that has a symmetric block-low-rank structure. In spirit, this layer is similar to bottleneck and squeeze-and-expand structures. We contribute symmetry by construction, and a combination of notation and flattening of tensors allows us to interpret these network structures in linear algebraic fashion as a block-low-rank matrix in factorized form and observe various properties. A video segmentation example shows that we can train a network to segment the entire video in one go, which would not be possible, in terms of memory requirements, using non-reversible networks and previously proposed reversible networks.

연구 동기 및 목표

• 3D 데이터 처리 과정에서 가역적 다운샘플링에 따른 채널 수 증가로 인한 가역 신경망의 메모리 폭발 문제를 해결한다.
• 가역성을 유지하면서도 컨볼루션 커널의 저장 용량과 계산량을 줄이는 파라미터 효율적인 레이어 구조를 개발한다.
• 기존 표준 가역 신경망의 메모리 한계로 인해 이전에는 불가능했던 전체 3D 영상 입력에 대한 완전히 가역 신경망의 엔드 투 엔드 학습을 가능하게 한다.
• 학습 중에 커널 랭크와 메모리 사용량을 명시적으로 제어할 수 있는 대칭 BLR 레이어를 제안한다.
• 학습 과정에서 단계적으로 레이어의 랭크를 증가시켜 표현력을 향상시키되, 메모리 사용량을 초기부터 증가시키지 않는 학습 가능한 적응형 랭크 전략을 설계한다.

제안 방법

• 컨볼루션 커널 행렬을 $ \mathbf{K}^\top \mathbf{K} $ 형태로 인수분해하는 대칭 블록-저랭크(BLR) 레이어를 도입하여, 구조적으로 양의 준정부성과 대칭성을 보장한다.
• 저랭크 구조를 가진 블록 행렬을 사용하여 선형 변환으로 표현함으로써, $ \mathbf{K}^\top \mathbf{K} $ 인수분해를 통해 효율적인 저장 및 계산을 가능하게 한다.
• 텐서 평탄화와 블록 행렬 표기법을 활용하여, 선형 대수학과 자동 미분에 호환되는 방식으로 레이어를 구조화된 선형 연산자로 간주한다.
• BLR 레이어를 완전히 가역적인 다수준 네트워크에 통합하여, 다운샘플링 및 업샘플링 단계에서 모두 가역성을 유지한다.
• 초기에는 저랭크 레이어로 시작하고 학습 과정에서 단계적으로 컨볼루션 커널을 추가하는 적응형 랭크 학습 전략을 구현하여 메모리 사용과 성능 간의 균형을 이룬다.
• BLR 구조를 활용하여 비선형 활성화 함수 적용 시에도 대칭성과 저랭크 구조 등의 수학적 성질을 관찰하고 유지한다.

실험 결과

연구 질문

• RQ1 대칭 블록-저랭크 레이어는 완전히 가역적인 다수준 네트워크에서 컨볼루션 커널의 메모리 용량을 줄이되, 가역성이나 표현력에 손상 없이 효과적으로 작용할 수 있는가?
• RQ2 기존의 파라미터 효율적인 설계(예: 1x1 컨볼루션, 블록 순환 커널, 스위즈-앤드-익스펜드 볼트넥)와 비교할 때 BLR 레이어의 메모리 사용량과 학습 효율성은 어떠한가?
• RQ3 BLR 레이어를 갖춘 완전히 가역 신경망은 기존 표준 가역 신경망의 메모리 한계로 인해 이전에는 불가능했던 전체 3D 영상 입력에 대해 엔드 투 엔드 학습이 가능한가?
• RQ4 적응형 랭크 학습 전략은 학습 전반에 걸쳐 낮은 메모리 사용량을 유지하면서도 고품질 모델로 수렴하는 데 기여하는가?
• RQ5 비선형 활성화 함수 적용 시 BLR 레이어의 어떤 수학적 성질이 유지되며, 이는 안정적인 학습을 어떻게 지원하는가?

주요 결과

• 대칭 BLR 레이어는 컨볼루션 커널의 메모리 용량을 $ O(n_{\text{chan}}^2) $ 에서 $ O(m n_{\text{chan}}) $ 로 줄여, 기존 가역 신경망에서 제한을 초래하는 제곱수 증가를 피한다. 여기서 $ m $ 은 블록-랭크이다.
• 이 방법은 이전에는 커널 저장 메모리가 다수의 다운샘플링 단계 이후에 폭발하여 불가능했던 전체 3D 영상 입력에 대한 완전히 가역 신경망의 학습을 가능하게 한다.
• BLR 레이어 구조는 기존의 가역 신경망 프레임워크에 직접 통합 가능하며, 맞춤형 행렬-벡터 곱 연산 또는 코드 변경 없이도 작동한다.
• 수치 실험을 통해 적응형 랭크 학습 전략이 영상 데이터에서 양호한 세그멘테이션 성능을 달성하며, 더 높은 메모리 요구량을 가진 모델의 성능을 따라잡는다.
• BLR 레이어의 대칭성과 양의 준정부성 구조는 안정성을 보장하며, 비선형 변환 하에서 레이어 성질에 대한 이론적 분석을 가능하게 한다.