Fast and Memory-Efficient Wavelet Convolutions via I/O-Aware Reformulation

2026-08-11

AI Summary

This work addresses the execution inefficiency and high memory footprint of Wavelet convolution, which are primarily constrained by high-bandwidth memory (HBM) data movement bottlenecks. For the first time, the operator is systematically optimized from an I/O perspective through three algebraic restructuring strategies: on-chip recomputation of Haar analysis butterfly operations, merging multi-level synthesis stages into a single closed-form computation, and fusing channel scaling factors directly into convolutional weights. These optimizations substantially reduce memory access overhead while preserving the exponential receptive field and linear parameter complexity inherent to Wavelet convolution. Experimental results demonstrate that, compared to the original implementation, the proposed approach reduces HBM traffic by 2.55×, accelerates training by up to 4.35×, and cuts peak memory consumption by nearly 50%.

Abstract

Wavelet convolution (WTConv) has emerged as an increasingly popular drop-in replacement for standard convolutions, expanding a network's receptive field exponentially with the number of decomposition levels while keeping the parameter count linear. However, its reference implementation is severely memory-bound due to excessive data movement through high-bandwidth memory (HBM). We develop an I/O model of WTConv to characterize this bottleneck and use it to guide three algebraic reformulations: (1) recomputing the inexpensive Haar analysis butterfly on chip, (2) collapsing the multi-level synthesis cascade into a single closed-form pass indexed by output-coordinate bits, and (3) folding learned per-channel scales into the convolution weights. Together, these reformulations enable an I/O-aware fused implementation that substantially reduces HBM traffic. We evaluate the WTConvNeXt configuration across decomposition levels and a broad range of tensor shapes. Despite performing comparable arithmetic, the reference WTConv is substantially slower than the depthwise convolution it replaces. Our reformulation reduces modeled HBM traffic by approximately $2.55\times$, yielding up to a $4.35\times$ training speedup over the reference while roughly halving peak memory usage. Thus, our reformulation preserves the benefits of WTConv while substantially reducing its execution time and memory footprint, removing the systems overhead that previously limited its practical efficiency.

Problem

• Wavelet convolution
• memory-bound
• high-bandwidth memory
• I/O bottleneck
• computational efficiency

Innovation

• Wavelet Convolution
• I/O-aware Reformulation
• Memory Efficiency
• Haar Transform
• Fused Implementation