Amit Aflalo, Shahaf E. Finder, Roy Amoyal, Eran Treister, Oren Freifeld
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.
Amit Aflalo, Eran Treister, Chaim Baskin, Oren Freifeld
Diffeomorphisms provide a flexible, topology-preserving mathematical tool for modeling complex spatial or temporal deformations and are used in various scientific fields. However, making these deformations simultaneously highly expressive and computationally efficient remains challenging. Continuous Piecewise-Affine Based (CPAB) transformations (Freifeld et al., ICCV 2015; IEEE TPAMI 2017) offer an attractive solution by parameterizing a family of diffeomorphisms via continuous velocity fields that are piecewise affine with respect to a chosen tessellation of the domain. Importantly, $d=\dim(θ)$, the dimension of the parameter vector $θ$, depends on the fineness of the tessellation rather than on the data resolution. Despite this compact parameterization, CPAB optimization has long been hindered by the tight coupling between trajectory integration and gradient computation with respect to $θ$. We introduce FG-CPAB, a factorized-gradient formulation for scalable CPAB transformations. By separating integration from parameter-space projection, our method reduces gradient computation complexity from $\mathcal{O}(dTN)$ to $\mathcal{O}(TN+dC)$, where $T$ is the number of integration steps, $N$ is the number of transformed points, and $C$ is the number of tessellation cells. This reformulation yields speedups of up to four orders of magnitude (e.g., $\approx 10^4\times$ in 2D), substantial memory savings (e.g., $\approx 500\times$ reduction in peak GPU memory usage in 2D), and dramatically higher expressiveness. This makes fine-tessellation CPAB practical in 2D and 3D for the first time, enabling highly expressive parametric diffeomorphisms.