[論文レビュー] WaveMix: A Resource-efficient Neural Network for Image Analysis

Pranav Jeevan, Kavitha Viswanathan+2 more|arXiv (Cornell University)|May 28, 2022

Advanced Neural Network Applications被引用数 11

ひとこと要約

WaveMixは、スケール不変性、シフト不変性、エッジのスパarsityといった画像の事前知識を活用するため、2次元離散ウェーブレット変換（2D-DWT）を用いたリソース効率の良いニューラルネットワークを提案する。損失なし、パラメータフリーのウェーブレットベースの特徴マッピングと解像度を維持するブロックを採用することで、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.

研究の動機と目的

• 多様なビジョンタスクに一般化できるが、計算コストとメモリコストを最小限に抑えたニューラルネットワークアーキテクチャの開発。
• ウェーブレット変換を用いて、自然画像に内在する事前知識（スケール不変性、シフト不変性、エッジのスパarsity）を包括的に活用すること。
• トランスフォーマーの硬直性や標準CNNの制限を避ける、柔軟でスケーラブルかつ解像度を維持するアーキテクチャの構築。
• ウェーブレットベースの特徴マッピングが、学習可能なフィルターやランダムフィルターを上回る精度と効率を達成できることの証明。
• アーキテクチャの大幅な見直しを伴わずに、セグメンテーションやスーパーキャラクタリゼーションなどのピクセル単位のタスクへの効率的適応を可能にすること。

提案手法

• WaveMixは、自己相似性と解像度を維持するブロックを採用し、特徴マップを多段階2次元離散ウェーブレット変換（2D-DWT）で処理することで、損失なしの空間的ダウンサンプリングとトークンマッピングを実現する。
• 2D-DWT操作は、画像の事前知識に基づいて空間情報を再編成し、パラメータ追加なしに特徴マップのサイズを縮小するとともに、受容野の成長を加速する。
• 各WaveMixブロックは、2D-DWTのローパスサブバンド出力に対して全結合層を用いてチャネルマッピングを実行し、効率的な長距離相互作用を可能にする。
• マッピング後に逆2次元離散ウェーブレット変換（逆2D-DWT）またはアップサンプリングを用いて元の空間解像度を回復させ、2次元構造を全体にわたって保持する。
• ウェーブレットフィルタは学習可能ではなく固定（例：ハール）であり、パラメータ効率性と安定性を確保する。アブレーションスタディにより、ランダムまたはフーリエベースの代替手法よりも優れていることが確認された。
• フレームワークは完全に畳み込み型であり、アーキテクチャの変更なしに可変な入力サイズとタスクをサポートする。トランスフォーマーやU-Netスタイルのモデルとは異なり、これにより柔軟性が向上する。

実験結果

リサーチクエスチョン

• RQ1 ウェーブレットベースのアーキテクチャは、パラメータ数とGPUメモリ使用量を減らしながらも、CNN やビジョントランスフォーマーを上回る精度を達成できるか？
• RQ2 2次元離散ウェーブレット変換（2D-DWT）は、深層学習においてスケール不変性、シフト不変性、エッジのスパarsityといった画像の事前知識をどれほど効果的に活用できるか？
• RQ3 2D-DWTを2次元フーリエ変換、ランダムフィルター、またはマックスプーリングに置き換えると、性能が低下し、リソース使用量が増加するか？
• RQ4 2D-DWT、チャネルマッピング、アップサンプリングのレイヤーの配置と順序が、モデルの性能と効率にどのように影響するか？
• RQ5 WaveMixは、アーキテクチャの再設計なしに、分類、セグメンテーション、スーパーキャラクタリゼーションといった多様なビジョンタスクに一般化可能か？

主な発見

• Cityscapes検証セットにおいて、63Mパラメータで82.70 mIoUを達成し、以前のSOTA（85Mパラメータで82.40 mIoU）を上回った。
• Places-365では、28Mパラメータで56.45%の精度を達成し、以前のSOTA（31Mパラメータで56.32%）を上回った。
• EMNIST Lettersでは、4Mパラメータで95.96%の精度を達成し、パラメータ数は同じである以前のSOTAと同等の性能を示した。
• EMNIST Digitsでは、4Mパラメータで99.82%の精度を達成し、以前のSOTAの307Mパラメータを大幅に削減した。
• 2D-DWTレイヤーを削除した場合、GPUメモリ使用量が62%削減された。これは、WaveMixのメモリ効率性における2D-DWTの貢献を裏付けた。
• 2D-DWTを2次元離散フーリエ変換（2D-DFT）に置き換えると、GPUメモリ使用量が73%増加し、精度が12%低下した。これにより、ウェーブレットベースのダウンサンプリングの優位性が明確になった。