arXiv:2608.10805cs.CVcs.AI2026-08

优化小波卷积的内存与速度瓶颈,实现更快更省显存的计算。

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

  • 通过重构算法减少数据在高带宽内存间的频繁移动。
  • 显存占用减半,训练速度最高提升4.35倍。
  • 适合追求高效部署的小波卷积应用开发者。

小波卷积(WTConv)作为标准卷积的即插即用替代方案,能以线性参数量实现指数级感受野扩展,但其参考实现严重受限于高带宽内存(HBM)的数据移动开销。本文建立WTConv的输入输出(I/O)模型以刻画该瓶颈,并据此提出三种代数重构:(1)在芯片上重算低成本的Haar分析蝴蝶结构;(2)将多层合成级联合并为一个由输出坐标位索引的闭式通路;(3)将每通道学习到的缩放因子融合进卷积权重。上述重构协同支持一种面向I/O的融合实现,显著降低HBM流量。我们在不同分解层级和多种张量形状下评估了WTConvNeXt配置。尽管运算量相近,参考实现仍远慢于其替代的深度可分离卷积。本方法使建模的HBM流量减少约2.55倍,训练速度最高提升4.35倍,同时峰值显存使用量减半。因此,新方法在保留小波卷积优势的同时,大幅降低执行时间与内存开销,消除了此前限制其实际效率的系统开销。

原文摘要 · Abstract (English)

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.

小波卷积内存优化加速计算硬件感知

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。