[논문 리뷰] WaveMix: A Resource-efficient Neural Network for Image Analysis

Pranav Jeevan, Kavitha Viswanathan외 2명|arXiv (Cornell University)|2022. 05. 28.

Advanced Neural Network Applications인용 수 11

한 줄 요약

WaveMix는 다중 수준 2차원 이산 웨이블릿 변환(2D-DWT)을 활용하여 척도 불변성, 이동 불변성, 에지 희박성 등의 영상 사전 지식을 효과적으로 활용하는 자원 효율적인 신경망 아키텍처를 제안한다. 손실 없는, 파rameter가 없는 웨이블릿 기반 특징 혼합과 해상도 유지 블록을 통해 WaveMix는 CNN, ViT 및 토큰 믹서보다 훨씬 적은 파라미터와 낮은 GPU 메모리 사용량으로 이미지 분류 및 의미 세그먼테이션 작업에서 최신 기술 수준의 정확도를 달성한다.

ABSTRACT

We propose a novel neural architecture for computer vision -- WaveMix -- that is resource-efficient and yet generalizable and scalable. While using fewer trainable parameters, GPU RAM, and computations, WaveMix networks achieve comparable or better accuracy than the state-of-the-art convolutional neural networks, vision transformers, and token mixers for several tasks. This efficiency can translate to savings in time, cost, and energy. To achieve these gains we used multi-level two-dimensional discrete wavelet transform (2D-DWT) in WaveMix blocks, which has the following advantages: (1) It reorganizes spatial information based on three strong image priors -- scale-invariance, shift-invariance, and sparseness of edges -- (2) in a lossless manner without adding parameters, (3) while also reducing the spatial sizes of feature maps, which reduces the memory and time required for forward and backward passes, and (4) expanding the receptive field faster than convolutions do. The whole architecture is a stack of self-similar and resolution-preserving WaveMix blocks, which allows architectural flexibility for various tasks and levels of resource availability. WaveMix establishes new benchmarks for segmentation on Cityscapes; and for classification on Galaxy 10 DECals, Places-365, five EMNIST datasets, and iNAT-mini and performs competitively on other benchmarks. Our code and trained models are publicly available.

연구 동기 및 목표

• 다양한 비전 작업에 대해 일반화 능력이 뛰어나면서도 계산 비용과 메모리 비용을 최소화하는 신경망 아키텍처를 개발하기 위해.
• 딥 러닝에서 웨이블릿 변환을 활용하여 천연 영상 사전 지식—척도 불변성, 이동 불변성, 에지 희박성—을 종합적으로 활용하기 위해.
• 변형기와 표준 CNN의 제약을 피하고 유연하고 확장 가능하며 해상도를 유지하는 아키텍처를 설계하기 위해.
• 웨이블릿 기반 특징 혼합이 학습 가능한 또는 무작위 필터보다 정확도와 효율성 측면에서 뛰어나다는 것을 입증하기 위해.
• 아키텍처의 대대적 개편 없이도 픽셀 단위 작업—세그먼테이션, 슈퍼레졸루션 등—에 효율적으로 적용 가능하도록 하기 위해.

제안 방법

• WaveMix는 자가 유사성과 해상도 유지 블록을 사용하여 2D-DWT를 활용해 특징 맵을 손실 없는 공간적 다운샘플링 및 토큰 혼합으로 처리한다.
• 2D-DWT 연산은 영상 사전 지식에 기반해 공간 정보를 재구성하여, 추가 파라미터 없이 특징 맵 크기를 줄이고 수신 영역 성장 속도를 높인다.
• 각 WaveMix 블록은 2D-DWT의 저통과 하위대역 출력에 대해 완전 연결층을 통해 채널 혼합을 적용하여 장거리 상호작용을 효율적으로 구현한다.
• 혼합 후 원래 공간 해상도를 복원하기 위해 역방향 2D-DWT 또는 업샘플링을 사용하여 전체 과정에서 2차원 구조를 유지한다.
• 웨이블릿 필터는 학습 가능한 것이 아니라 고정되어 있음(예: 하르)으로 파라미터 효율성과 안정성을 확보하며, 추론 실험을 통해 무작위 또는 푸리에 기반 대안보다 뛰어남을 확인하였다.
• 이 프레임워크는 완전 컨볼루션 구조이며, 아키텍처 변경 없이도 다양한 입력 크기와 작업을 지원한다. 이는 변형기나 U-Net 스타일 모델과는 다릅니다.

실험 결과

연구 질문

• RQ1 웨이블릿 기반 아키텍처가 파라미터와 GPU 메모리 사용량을 줄이며도 CNN 및 비전 트랜스포머를 능가하는 정확도를 달성할 수 있는가?
• RQ2 2D-DWT가 딥 러닝에서 척도 불변성, 이동 불변성, 에지 희박성 등의 영상 사전 지식을 얼마나 효과적으로 활용하는가?
• RQ3 2D-DWT를 2D-DFT, 무작위 필터 또는 최대 풀링으로 대체할 경우 성능 저하와 자원 사용 증가가 발생하는가?
• RQ4 2D-DWT, 채널 혼합, 업샘플링 레이어의 배치 및 순서가 모델 성능과 효율성에 어떤 영향을 미치는가?
• RQ5 WaveMix는 아키텍처 재설계 없이도 분류, 세그먼테이션, 슈퍼레졸루션 등 다양한 비전 작업에 일반화 가능한가?

주요 결과

• Cityscapes 검증 세트에서 WaveMix는 단 63M 파라미터로 82.70 mIoU를 달성하여 이전 최고 성능(85M 파라미터에서 82.40 mIoU)을 초월했다.
• Places-365에서 WaveMix는 28M 파라미터로 56.45%의 정확도를 기록하여 이전 최고 성능(31M 파라미터에서 56.32%)을 능가했다.
• EMNIST Letters에서 WaveMix는 단 4M 파라미터로 95.96%의 정확도를 달성하여 이전 최고 성능와 동일한 파라미터 수를 유지하면서도 동일한 성능을 달성했다.
• EMNIST Digits에서 WaveMix는 4M 파라미터로 99.82%의 정확도를 기록하여 이전 최고 성능의 307M 파라미터 수치를 크게 감소시켰다.
• 2D-DWT 레이어를 제거했을 때 GPU RAM 사용량이 62% 감소하여, 이는 WaveMix의 메모리 효율성에 기여하는 바가 크다는 것을 확인했다.
• 2D-DWT를 2D-DFT로 대체했을 경우 GPU RAM 사용량이 73% 증가하고 정확도가 12% 감소하여 웨이블릿 기반 다운샘플링의 우수성을 입증했다.