Fast and Memory-Efficient Wavelet Convolutions via I/O-Aware Reformulation
Amit Aflalo, Shahaf E. Finder, Roy Amoyal, Eran Treister, Oren Freifeld
Ben-Gurion University of the Negev
cs.CV, cs.AI
Submitted: 2026-08-13
Updated: 2026-08-14
License: http://creativecommons.org/licenses/by/4.0/
Importance score: 75/100
The gist: 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
Terminology
Summary
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-up iterations. Baselines: WTConv (reference and fused) runs at k=5 throughout § 5.2 and § 5.3; § 5.4 repeats the sweep at k=3, and § 5.7 repeats it at k=5 on a second GPU. We evaluate plain depthwise convolutions under two protocols: matched (k=5) to isolate wavelet overhead, and drop-in (k=7) representing the standard ConvNeXt replacement. Reporting: Table values are the geometric mean over the (C, H) sweep of the per-configuration ratio between the WTConv reference and the evaluated method. The reference is always 1.00×. Latency is reported as speedup (reference/method; higher is better) and memory as footprint fraction (method/reference; lower is better).
For a full training step, the reference operator is slower than the convolution it replaces: the depthwise 7 × 7 convolution of a ConvNeXt block completes a training step 2.46–3.42× faster in fp32 and 1.53–2.20× faster in fp16; at matched kernel size the margin widens to 4.89–6.79× (fp32, depthwise 5 × 5). Fusion reverses this. The fused layer is 3.71–4.35× faster than the reference in fp32 and 2.68–3.09× in fp16. The ratio grows with L and then flattens, because deeper decompositions give the reference more intermediates to materialize while Eq. 10 is nearly flat in L; the change from L=4 to L=5 is below 1% in both precisions, as Eq. 11 predicts. At every level and in both precisions this suffices to overturn the comparison above: against the depthwise 7 × 7 convolution it replaces, the fused layer trains 1.27–1.50× faster in fp32 and 1.40–1.76× faster in fp16. It does not win against a depthwise convolution at its own kernel size, which remains 1.32–1.57× faster in fp32 and 2.34–2.94× in fp16; that kernel moves 2N elements against the fused layer’s 7N –8.33N and carries no decomposition at all, so the direction is expected. The speedup over the reference is smaller in fp16 than in fp32 because half precision halves the bytes moved by both implementations, and the reference, which moves more of them, benefits more.
For forward-only latency, against the reference the picture is unchanged: the fused layer is 3.67–4.37× faster in fp32 and 3.38–3.60× in fp16, and the reference remains slower than both plain convolutions in the table, including the depthwise 7×7 it is meant to replace, by 3.79–5.28× and 2.05–2.70×. The comparison against the baselines, however, does not reverse as cleanly as over a training step. Against the depthwise 7 × 7 convolution the fused layer wins in fp16 (1.31–1.65×) but not in fp32, where it sits at 0.83–0.97×.
For peak allocated memory over a training step, the fused layer needs 0.43–0.55× the memory of the reference (a reduction by a factor of 1.83–2.31), and the reduction is essentially flat in L beyond L=2, because the dominant saved allocation is the level-1 coefficient tensor: it holds N elements, three times as many as all deeper levels combined (Σ l≥2 Nl → N/3). This behavior follows from the mechanism described in § 4: every per-level subband tensor the reference writes to HBM is also retained by autograd until the backward pass, whereas the fused layer recomputes the Haar coefficients on chip. Against the plain convolution, the fused layer’s training-step footprint is 0.87–1.02× that of the depthwise 7 × 7 convolution in fp32, and lower still in fp16 (0.76–0.89×). The fused L-level layer therefore has a training footprint within a few percent of the memory budget of a single plain convolution and falls below it in half precision.
For forward-only inference, the fused/reference footprint ratio is 0.50–0.65× during inference, versus 0.43–0.55× during training. The smaller reduction at inference is expected because a large part of what fusion removes is tape-retained subband tensors that inference never allocates in the first place; what remains is the traffic-side reduction of Eq. 10. The plain convolutions are correspondingly harder to match, since without a tape their footprint is close to the compulsory input plus output: the fused layer sits at 1.15–1.34× the depthwise 7× 7 convolution’s footprint in fp32 and 0.99–1.15× in fp16, the residual being the filtered subband coefficients it must materialize between the analysis and synthesis passes.
Every measurement above runs at k=5, the kernel size WTConvNeXt deploys. The modeled tensor-materialization traffic, however, is independent of k: Eq. 2 and Eq. 10 are both independent of k, which sets how much arithmetic each resident element receives rather than how many elements cross the HBM boundary. Table 6 tests that claim by repeating the sweep of § 5.1 at k=3. The speedup over the reference is larger at k=3 than at k=5 in every cell of the table. A training step is 3.86–4.68× faster than the reference in fp32 and 3.31–3.94× in fp16, against 3.71–4.35× and 2.68–3.09× at k=5; inference is 4.48–5.32× and 4.40–4.84×, against 3.67–4.37× and 3.38–3.60×. § 3 predicts this direction. Shrinking the kernel removes arithmetic while leaving Eq. 2 untouched, so the reference falls further into the memory-bound regime.
To evaluate the operator at architectural scale, we benchmark ConvNeXt-T and WTConvNeXt-T, which replaces ConvNeXt-T’s depthwise convolutions with WTConv (k=5, L∈ 5, 4, 3, 2). We use a batch size of 64, 224 × 224 inputs, and fp32. Because the reformulation preserves the mathematical operator up to floating-point evaluation order, it does not alter the model architecture or learned parameters. The reference WTConv imposes a severe bottleneck, restricting the network to 30% of ConvNeXt-T’s inference throughput and 67% of its training throughput. Our fused implementation effectively mitigates this, achieving 2.35× and 1.58× the reference’s inference and training throughputs, respectively. Crucially, WTConvNeXt-T equipped with the fused layer trains 1.06× faster than the baseline ConvNeXt-T, although inference throughput reaches only 0.71× of the baseline. The fused formulation similarly reduces peak memory, cutting the reference WTConvNeXt-T’s training and inference footprints to 0.79× and 0.85×. This limits the network’s peak memory to 1.10× that of ConvNeXt-T during training and 1.26× during inference.
We evaluate the three reformulations cumulatively, in their order in § 4. The first variant fuses Haar analysis, depthwise convolution, and scaling at each level. The second also replaces the sequential synthesis loop with the bit-indexed pass of Proposition 2. The full method further folds the base-path scale and fuses its addition into the final store. Each variant otherwise follows the reference WTConv2d implementation and is verified against its outputs and parameter gradients. We measure full training steps at k=5 in fp32, using the (C, H) sweep of § 5.1. The training-step speedup relative to the reference (1.00×) is: + Fused Haar analysis: 1.35×, 1.44×, 1.46×, 1.45×, 1.44× for L=1..5; + Collapsed synthesis: 2.54×, 2.95×, 3.02×, 3.04×, 3.05×; + Scale folding = Full: 3.71×, 4.23×, 4.34×, 4.35×, 4.33×. The peak training-step memory relative to the reference (1.00×; lower is better) is: + Fused Haar analysis: 0.79×, 0.60×, 0.59×, 0.59×, 0.59×; + Collapsed synthesis: 0.65×, 0.49×, 0.50×, 0.50×, 0.50×; + Scale folding = Full: 0.55×, 0.43×, 0.44×, 0.45×, 0.45×.
All preceding layer measurements use an RTX A6000 (Ampere, sm 86). To assess device dependence, Table 10 repeats the k=5 sweep on an NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition (sm 120). The fused implementation remains faster for all L and both precisions. Relative to the reference, training is 2.46–2.79× faster in fp32 and 1.95–2.42× faster in fp16; inference is 2.87–3.05× and 2.92–3.10× faster, respectively. The realized speedup is device-dependent even though both implementations remain well below the device’s roofline ridge point. Differences in achieved bandwidth, kernel dispatch, launch overhead, cache behavior, and kernel utilization are not captured by the tensor-traffic model. Peak-memory reductions are stable across devices, since they are governed mainly by tensor materialization rather than GPU throughput, matching the A6000 trends at 0.40–0.55× for training and 0.50–0.65× for inference.
Our reformulation relies on properties specific to the Haar wavelet, including its low-arithmetic-cost, symmetric, and self-inverse transform structure. In particular, the fused analysis and bit-indexed closed-form synthesis do not directly generalize to wavelet families with longer filters or more complex reconstruction rules. Extending the approach beyond Haar would therefore require new transform-specific formulations.
WTConv’s performance bottleneck is not arithmetic but data movement. Although the operator performs only moderately more computation than the depthwise convolution it replaces, its reference implementation repeatedly materializes intermediate wavelet coefficients and reconstructions in HBM, leaving it deeply memory-bound. By making this I/O cost explicit, we derived an algebraically equivalent formulation that keeps Haar analysis on chip, collapses the multi-level synthesis recursion into a single bit-indexed pass, and folds learned scales into the convolution weights. These reformulations reduce modeled HBM traffic by approximately 2.55× and translate directly into substantial practical gains. Across the evaluated configurations, the fused implementation is 3.71–4.35× faster than the reference in fp32 and 2.68–3.09× faster in fp16 over a full training step, while reducing peak memory by a factor of 1.83–2.31. More importantly, the optimization changes the practical trade-off that motivates WTConv: the fused layer trains 1.27–1.50× faster in fp32 and 1.40–1.76× faster in fp16 than the depthwise 7 × 7 convolution it is intended to replace. The broader lesson is that favorable FLOP counts and parameter scaling do not by themselves imply an efficient operator. For structured, multi-stage layers built from inexpensive transforms, intermediate tensor materialization can dominate execution cost. In such settings, I/O-aware algebraic reformulation is not merely an implementation optimization; it can determine whether the theoretical advantages of an operator translate into practical gains.
Improvements for AI systems
Improvements to AI systems:
-
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.
-
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.
-
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.
-
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.
-
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.
-
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.
-
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.
-
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.
-
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).
-
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.
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.
Related papers
- Loss Knows Best: Detecting Annotation Errors in Videos via Loss Trajectories
- AnchorWeave: World-Consistent Video Generation with Retrieved Local Spatial Memories
- Benchmarking the Robustness of Foundation Models for Mammography under Domain Shift
- MambaX-Net: Dual-Input Mamba-Enhanced Cross-Attention Network for Longitudinal MRI Segmentation
- TeleOCR: Navigating Document Parsing Across Digital and Camera-Captured Documents
- A Survey on Efficient Vision-Language-Action Models