[论文解读] 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 提出了一种用于图像分析的资源高效神经网络，通过利用多级二维离散小波变换（2D-DWT）来挖掘图像先验特性，如尺度不变性、平移不变性和边缘稀疏性。通过使用无损、无参数的小波特征混合与保持分辨率的模块，WaveMix 在图像分类和语义分割任务上实现了最先进的准确率，同时参数量显著减少，GPU 显存使用量也低于 CNN、ViT 和 token mixer 模型。

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.

研究动机与目标

• 开发一种在多种视觉任务中具有良好泛化能力的神经网络架构，同时将计算和内存成本降至最低。
• 在深度学习中全面利用自然图像先验特性——尺度不变性、平移不变性和边缘稀疏性，通过小波变换实现。
• 构建一种灵活、可扩展且保持分辨率的架构，避免 Transformer 的僵化性与标准 CNN 的局限性。
• 证明基于小波的特征混合在准确率和效率方面优于可学习或随机滤波器。
• 在无需架构重构的前提下，实现对像素级任务（如分割和超分辨率）的高效适配。

提出的方法

• WaveMix 使用自相似、保持分辨率的模块，通过多级 2D-DWT 处理特征图，实现无损空间下采样与 token 混合。
• 2D-DWT 操作基于图像先验重新组织空间信息，不增加参数即可减小特征图尺寸，并加速感受野增长。
• 每个 WaveMix 模块在 2D-DWT 的低通子带输出上通过全连接层进行通道混合，实现高效的长距离交互。
• 通过逆 2D-DWT 或上采样操作在混合后恢复原始空间分辨率，全程保持 2D 结构。
• 小波滤波器为固定值（例如 Haar），而非可学习参数，确保参数效率与稳定性，消融实验确认其优于随机或基于傅里叶的替代方案。
• 该框架为全卷积结构，支持可变输入尺寸与任务，无需架构更改，与 Transformer 或 U-Net 类模型不同。

实验结果

研究问题

• RQ1 基于小波的架构是否能在参数更少、GPU 显存使用更少的情况下，超越 CNN 和视觉 Transformer 的准确率？
• RQ2 2D-DWT 在深度学习中多大程度上有效利用了图像先验特性，如尺度不变性、平移不变性和边缘稀疏性？
• RQ3 若将 2D-DWT 替换为傅里叶变换、随机滤波器或最大池化，是否会导致性能下降与资源使用增加？
• RQ4 2D-DWT、通道混合与上采样层的排列顺序与位置如何影响模型性能与效率？
• RQ5 WaveMix 是否能在无需架构重设计的前提下，泛化于多种视觉任务（包括分类、分割与超分辨率）？”

主要发现

• 在 Cityscapes 验证集上，WaveMix 达到 82.70 mIoU，仅使用 63M 参数，优于之前的 SOTA（82.40 mIoU，85M 参数）。
• 在 Places-365 上，WaveMix 达到 56.45% 准确率，仅使用 28M 参数，优于先前 SOTA（56.32%，31M 参数）。
• 在 EMNIST Letters 上，WaveMix 达到 95.96% 准确率，仅使用 4M 参数，与先前 SOTA 相当，且参数量相同。
• 在 EMNIST Digits 上，WaveMix 达到 99.82% 准确率，仅使用 4M 参数，显著低于先前 SOTA 的 307M 参数量。
• 当移除 2D-DWT 层后，WaveMix 将 GPU 显存使用量降低 62%，证实其在内存效率中的关键作用。
• 将 2D-DWT 替换为 2D-DFT 后，GPU 显存使用量增加 73%，准确率下降 12%，证明基于小波的下采样更具优势。