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

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×, yielding up to a 4.35× 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. Source code is available in the official WTConv repository under fast wtconv.

Large receptive fields are important for modern convolutional networks, but obtaining them with conventional convolutions is expensive. Stacking small kernels expands the theoretical receptive field only linearly with depth, while directly increasing kernel size incurs parameter and arithmetic costs proportional to the kernel area. Fortunately, the widely-used WTConv offers an appealing alternative: it applies small depthwise convolutions across the progressively downsampled levels of a wavelet decomposition, so the receptive field grows exponentially with the number of levels while the parameter count grows only linearly. This makes WTConv a drop-in replacement for large depthwise convolutions. WTConv has also produced accuracy and robustness gains in architectures including ConvNeXt and MobileNetV2.

Yet this advantage does not translate into execution speed. In the original WTConv paper, WTConvNeXt uses WTConv with a 5 × 5 kernel (k=5) as a replacement for the 7 × 7 depthwise convolution in a ConvNeXt block. In this setting, the reference implementation is slower than the convolution it replaces: over a full training step, it trails the depthwise 7 × 7 baseline by 2.46–3.42× in fp32 and 1.53–2.20× in fp16; at inference, the corresponding gaps are 3.79–5.28× and 2.05–2.70×. These gaps cannot be explained by arithmetic alone. For an input containing N = B · C · H · W elements, WTConv with k=5 performs 58N –69N multiply–accumulates as the number of decomposition levels increases from L=1 to L=5, compared with 49N for the 7 × 7 depthwise baseline. This increase in arithmetic is far smaller than the observed latency gap. Thus, despite its favorable parameter scaling, the reference operator incurs a substantial wall-clock penalty. Closing this performance gap is the primary objective of this work.

The discrepancy arises because FLOPs are the wrong cost model for this operator. The reference WTConv implementation has an arithmetic intensity of only ≈ 1.63 FLOP/byte in fp32, placing it deep in the memory-bound regime on modern GPUs. Its execution is therefore dominated not by the cost of the Haar transform or the depthwise convolutions themselves, but by repeatedly materializing intermediate wavelet coefficients and reconstructions in high-bandwidth memory (HBM). Under a tensor-materialization I/O model, a forward evaluation incurs 17.75N –21.32N element reads and writes to global memory for an input containing N elements.

This observation suggests a different optimization target: rather than reducing arithmetic, we reformulate WTConv so that mathematical intermediates need not become memory-resident intermediates. We exploit three properties of the operator. First, Haar analysis uses only signed additions and fixed power-of-two scaling and can be recomputed on chip inside the depthwise convolution. Second, linearity of Haar synthesis allows the entire L-level reconstruction recursion to be written as a single closed-form sum whose signs and coefficient addresses are determined by bits of the output coordinate. Third, the learned per-channel scales can be folded into the convolution weights. Together, these transformations preserve the same mathematical operator while eliminating the large intermediate tensors responsible for the majority of its data movement. The resulting formulation reduces modeled HBM traffic by 2.54–2.56× for L = 1,..., 5.

We implement this formulation in CUDA and evaluate both inference and full training steps over a broad sweep of tensor shapes, decomposition levels, and precisions. For a full training step, the implementation is 3.71–4.35× faster than the reference in fp32 and 2.68–3.09× faster in fp16, while reducing peak memory by a factor of 1.83–2.31. More importantly, the systems reformulation reverses the practical comparison that motivates the work: over a training step the optimized WTConv at k=5 is 1.27–1.50× faster in fp32 and 1.40–1.76× faster in fp16 than the depthwise 7 × 7 convolution it is proposed to replace, at every decomposition level tested. Training is the regime in which the comparison matters most, since it is where the materialized intermediates are both written and re-traversed; at inference the reversal is complete in fp16 but not in fp32, where the fused layer remains within 3–17% of the depthwise baseline rather than ahead of it.

Contributions:

• An I/O cost model for WTConv. We derive an element-level accounting of HBM traffic for the reference implementation, Qref = 7N + 43/3 N (1 − 4−L), which is independent of the kernel size, and use a roofline analysis to show that WTConv’s arithmetic intensity is a factor of roughly 31 below the compute-saturation ridge point in fp32 at k=5. This identifies data movement, rather than arithmetic, as its dominant cost.

• An algebraically exact, I/O-aware reformulation. We combine register-resident Haar analysis, a closed-form bit-indexed synthesis over all decomposition levels, and scale folding to eliminate unnecessary HBM-resident intermediates. The resulting formulation reduces predicted forward-pass traffic by 2.54–2.56× across L = 1,..., 5 while preserving the WTConv operator up to floating-point evaluation order.

• End-to-end empirical validation. Across all measured configurations, our CUDA implementation substantially reduces training and inference latency and peak memory relative to the reference WTConv implementation, and over a training step outperforms the depthwise 7 × 7 convolution WTConv is proposed to replace. We additionally verify numerical agreement for the forward output and all parameter gradients.

The reference implementation expresses each level as a chain of framework primitives, every one of which reads its input from HBM and writes its output back. We count elements crossing the HBM boundary: Each stage reads its input tensor once and writes its output tensor once. Weight traffic is O(Ck2) and therefore negligible against O(N). Writing Nl = N/4(l−1) for the element count entering level l:

• Analysis, per level: the Haar transform is executed as a grouped stride-2 conv2d against a constant ±½ filter (2Nl); the depthwise convolution over the 4C-channel coefficient tensor (2Nl); the scale multiply (2Nl). Total 6Nl.

• Synthesis, per level: the cross-level low-pass add (4/3 Nl); the concatenation of the summed low-pass band with the three high bands (2Nl); the conv transpose2d (2Nl). Total 19/4 Nl.

• Base path: convolution (2N), scale multiply (2N), final add (3N). Total 7N.

Using Σ l=1 L 4(−(l−1)) = 4/3 (1 − 4(−L)), Qref = 7N + (6 + 19/4) · 4/3 N (1 − 4(−L)) = 7N + 43/3 N (1 − 4(−L)), rising from 17.75N at L = 1 to 21.33N as L → ∞; that is, evaluating the layer moves between 18 and 21 times the input tensor’s worth of data through HBM. Note that k does not appear: the kernel size sets how much arithmetic each resident element receives, not how many elements cross the boundary. For a representative B=8, C=64, H=W=256 input, N = 33.6M elements (128 MiB in fp32) and the reference therefore moves roughly 2.9 GB per forward pass.

The base convolution performs k2 multiply–accumulates per element. At level l, the depthwise convolution over the coefficient tensor costs a further k2 per element of Nl, and the analysis and synthesis each cost 4, since the reference executes both as 2 × 2 convolutions rather than as additions. Summing over levels with the same geometric factor as before, MAC = N [k2 + (k2 + 8) · 4/3 (1 − 4(−L))]. For k = 5 and L = 5 this is 69.0N multiply–accumulates, or 137.9N FLOPs, against 4Qref = 85.3N bytes in fp32: Iref = 137.9N / 85.3N ≈ 1.63 FLOP/byte. Against the ridge point I⋆ ≈ 50, WTConv’s arithmetic intensity is lower by a factor of roughly 31 in fp32; equivalently, it is approximately 3.2% of the ridge-point intensity. Enlarging the kernel does not change this conclusion, only its margin: increasing k increases the arithmetic work while leaving Eq. 2 untouched, so even at k=5 the operator remains more than an order of magnitude below the ridge point. Arithmetic is not the bottleneck: accelerating the transform itself would have little effect, whereas reducing memory traffic directly targets the dominant cost.

To eliminate the severe memory bottleneck of the reference WTConv, we derive an I/O-aware formulation based on three exact algebraic reformulations. These reformulations leave the underlying mathematical operator unchanged, differing from the reference implementation solely in floating-point evaluation order. Together, they eliminate the unnecessary Haar-analysis and reconstruction intermediates from high-bandwidth memory (HBM) by performing multi-stage computations directly in GPU registers and shared memory across all decomposition levels.

For a 2×2 spatial patch [[a, b], [c, d]], the normalized 2D Haar analysis transform maps four input pixels to subband coefficients (LL, LH, HL, HH) as follows: LL = ½(a + b + c + d), LH = ½(a + b − c − d), HL = ½(a − b + c − d), HH = ½(a − b − c + d). Evaluating these requires only signed additions and fixed scaling by ½. In contrast, the reference implementation materializes the entire coefficient tensor Y(l) in HBM via grouped convolutions with ±½ filters. Each level runs Haar analysis, depthwise convolution, and learned scaling as separate passes. Each pass reads and writes an Nl-element tensor, moving 6Nl elements in total. Our fused approach computes the Haar coefficients inside the depthwise convolution and folds the scale into its weights. Because WTConv reaches only 1.63 FLOP/byte, recomputing the coefficients is much cheaper than materializing them in HBM.

Naive fusion would recompute the transform for every convolution tap and load each input pixel k2 times. Instead, each thread block transforms the 2 × 2 blocks for one output tile, including a coefficient-space halo of radius R = (k − 1)/2, and stages the four subbands in shared memory. All convolution taps reuse the transformed tile, so each 2 × 2 input block is transformed only once per output tile. Except at the deepest level, the fused approach reads Nl input elements, writes Nl filtered coefficients for synthesis, and writes ¼Nl unfiltered low-pass coefficients for the next level. It therefore moves 9/4 Nl elements, versus 6Nl for the reference. The deepest level has no successor and moves 2NL.

The reference implementation executes reconstruction as an L-step sequential loop, where each level reads the lower-level reconstruction from HBM, adds it to its low-pass subband, and applies transposed convolution. This imposes L strict sequential dependencies and incurs an I/O cost of 19/4 Nl per level. Because Haar synthesis is linear, the low-pass carrier acts as a linear accumulator across levels. As a result, the multi-level synthesis cascade can be collapsed into a single closed-form pass across all L levels. Haar synthesis at level l expands each subband coefficient over a 2l × 2l pixel block using a fixed sign pattern determined by the spatial location of the output pixel within that block.

Proposition 2 (Bit-indexed single-pass synthesis). Let (y, x) denote an output pixel coordinate, and let level l ∈ 1,..., L have subband coefficients (cLL(l), cLH(l), cHL(l), cHH(l)) at coarse grid location (⌊y/2l⌋, ⌊x/2l⌋). Define the coordinate parity signs at level l as follows: s y(l) = (−1)(⌊y/2(l−1)⌋ mod 2), s x(l) = (−1)(⌊x/2(l−1)⌋ mod 2). The reconstructed pixel value Ry,x across all L levels is given by: Ry,x = Σ l=1 L 2(−l) [cLL(l) + s y(l) cLH(l) + s x(l) cHL(l) + s y(l) s x(l) cHH(l)]. By evaluating this in a single pass, a GPU thread computing pixel (y, x) iterates from level L down to 1, addressing coarse coefficients via coordinate bit-shifts and obtaining signs through bitwise parity checks. This eliminates all intermediate low-pass tensors and sequential kernel dependencies. Furthermore, the base-path convolution output is folded directly into the final output store, eliminating a separate full-resolution tensor addition.

The reference implementation applies learned per-channel scaling sc after convolution (Zc = sc (X ∗dw W)c) as a standalone elementwise pass. At an arithmetic intensity of 18 FLOP/byte in fp32 (one multiply per element read and written), this standalone pass is strongly memory-bound and adds an avoidable read–write traversal of the tensor. By the linearity of convolution, the scale factor is folded into the depthwise weights and, on the base path, into its bias: W̃c = sc ⊙ Wc, b̃c = sc bc, so that Z = X ∗dw W̃ + b̃ is computed in a single pass with zero extra runtime cost. During backpropagation, exact parameter gradients with respect to the unscaled weight Wc, bias bc, and scale sc are efficiently recovered as follows: ∂L/∂Wc = sc ∂L/∂W̃c, ∂L/∂bc = sc ∂L/∂b̃c, ∂L/∂sc = ∂L/∂W̃c · Wc + ∂L/∂b̃c · bc. sc enters both folds, so its gradient collects a term from each.

Collecting the reformulations across all levels 1,..., L, with Nl = N/4(l−1):

• Base path: the folded vendor convolution reads N and writes N (2N).

• Analysis & Depthwise pass, levels 1 ≤ l ≤ L: reads Nl, writes Nl convolved coefficients and ¼Nl raw low-pass coefficients (9/4 Nl each, less the ¼NL the deepest level does not write).

• Single-pass Synthesis: reads all subband coefficients Σ l=1 L Nl and base-path output N, writing the final output N (Σ l=1 L Nl + 2N).

Using Σ l=1 L Nl = 4/3 N (1 − 4(−L)), the total memory traffic is: Qfused = 4N + 13/3 N (1 − 4(−L)) − N 4(−L). Comparing this to the reference memory traffic Qref, the theoretical reduction ratio in HBM data movement is: Qref/Qfused = [7 + 43/3 (1 − 4(−L))] / [4 + 13/3 (1 − 4(−L)) − 4(−L)] → 64/25 = 2.56 as L → ∞. As L → ∞, the fused formulation achieves up to a 2.56× reduction in memory traffic over the reference implementation. Table 1 summarizes the exact element counts and traffic reduction ratios for decomposition levels L=1,..., 5. Like Eq. 2, neither count depends on k, so the traffic argument is stated once and holds at every kernel size; § 5.4 verifies this empirically at k=3.

All measurements are single-layer microbenchmarks on an RTX A6000, except the end-to-end networks in § 5.5 and the cross-hardware validation in § 5.7, which repeats the layer sweep on a second, architecturally different GPU. We sweep C ∈ 32, 64, 128, H = W ∈ 128, 256, 512 at B=8, and L ∈ 1,..., 5 in fp32 and fp16. Timings are averaged over 50 iterations following 20 warm…

Improvements to AI systems:

1. I/O-aware operator fusion for structured transforms: AI systems can automatically detect when multi-stage operators (wavelet transforms, FFTs, or other linear decompositions) materialize intermediate tensors in HBM, then apply algebraic reformulations—such as on-chip recomputation of cheap transforms, closed-form single-pass synthesis via bit-indexed coordinate addressing, and folding learned scales into weights—to reduce memory traffic by 2.55× and achieve 3.7–4.35× training speedups.

2. Memory-efficient training for large receptive field networks: AI systems using WTConv or similar wavelet-based architectures can now train 1.27–1.50× faster than standard depthwise 7×7 convolutions in fp32 (1.40–1.76× in fp16) while reducing peak memory by 1.83–2.31×, enabling deployment of exponentially-growing receptive fields on memory-constrained hardware without sacrificing training throughput.

3. Roofline-guided optimization: AI systems can use the derived I/O cost model (Qref = 7N + 43/3·N·(1−4−L)) to predict memory-bound bottlenecks before implementation, identifying when arithmetic intensity falls below the ridge point (here, 31× below) and prioritizing data-movement reduction over FLOP optimization for structured operators.

4. Exact gradient recovery under weight folding: AI systems can safely fold learned per-channel scales into convolution weights during forward passes while recovering exact parameter gradients (∂L/∂W = s·∂L/∂W̃, ∂L/∂s = ∂L/∂W̃·W + ∂L/∂b̃·b) without additional memory overhead, enabling faster training with mathematically identical updates.

5. Cross-device portable memory optimizations: AI systems can achieve stable peak-memory reductions (0.40–0.55× training, 0.50–0.65× inference) across different GPU architectures (Ampere, Blackwell) even when speedups vary, because memory savings derive from tensor-materialization elimination rather than hardware-specific tuning—making the approach robust for heterogeneous deployments.

6. Automatic synthesis-cascade collapsing: AI systems can replace recursive multi-level reconstruction loops with a single closed-form pass where output-coordinate bits determine both coefficient addresses and sign patterns (s y = (−1)(⌊y/2(l−1)⌋ mod 2)), eliminating sequential dependencies and intermediate low-pass tensors, reducing I/O from 19/4·Nl to a single traversal.

7. Precision-aware optimization: AI systems can automatically adjust fusion strategies based on precision—fp16 benefits more from traffic reduction (2.68–3.09× training speedup) than fp32 (3.71–4.35×) because half precision halves bytes moved for both implementations—allowing dynamic selection of optimization depth based on numeric format.

8. End-to-end architectural throughput improvements: AI systems can now deploy WTConvNeXt-T with fused layers to achieve 2.35× inference and 1.58× training throughput over reference implementations, and even train 1.06× faster than baseline ConvNeXt-T, making wavelet-based large-receptive-field models practical for production workloads.

9. Kernel-size-independent optimization: AI systems can apply the same I/O reformulation regardless of convolution kernel size (k=3, k=5), since memory traffic is independent of k—the optimization remains effective even as arithmetic varies, with speedups increasing at smaller kernels (3.86–4.68× vs 3.71–4.35× for training).

10. Backward-pass memory reduction via recomputation: AI systems can choose to recompute cheap transforms (like Haar analysis at 1.63 FLOP/byte) during backpropagation instead of storing intermediates, reducing autograd tape memory by up to 2.31× while adding negligible compute—a generalizable strategy for any low-arithmetic-intensity transform.

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.