文章背景与核心概要
小波卷积(WTConv)作为标准卷积的一种强大替代方案,能够以线性的参数扩展实现呈指数增长的感受野。然而,由于高带宽内存(HBM)中存在大量的数据移动,导致了严重的内存瓶颈,从而阻碍了其在实际中的广泛应用。
本文引入了一种感知I/O的模型来识别这些瓶颈,并提出了三种代数重构方法来优化性能:1. 片上Haar分析:直接在芯片上重新计算低成本的Haar分析蝶形运算;2. 闭式综合:将多级综合级联折叠为由输出坐标位索引的单次遍历;3. 权重折叠:将学习得到的每个通道的缩放比例直接整合到卷积权重中。与参考实现相比,这些优化使HBM流量减少了2.55倍,训练速度提升了4.35倍,峰值内存使用量减少了50%,从而有效消除了以往限制WTConv效率的系统级开销。
Fast and Memory-Efficient Wavelet Convolutions via I/O-Aware Reformulation
Authors: Amit Aflalo, Shahaf E. Finder, Roy Amoyal, Eran Treister, Oren Freifeld
arXiv: 2608.10805 [cs.CV]
Submitted: 11 Aug 2026 (v1), 13 Aug 2026 (v2)
Summary
小波卷积(WTConv)是标准卷积的一种强大替代方案,它以线性的参数扩展提供呈指数增长的感受野。然而,由于高带宽内存(HBM)中过多的数据移动导致了严重的内存瓶颈,其实际应用一直受到阻碍。
Wavelet convolution (WTConv) is a powerful alternative to standard convolutions, offering an exponentially increasing receptive field with linear parameter scaling. However, its practical adoption has been hindered by significant memory bottlenecks caused by excessive data movement through high-bandwidth memory (HBM).
本文引入了一种感知I/O的模型来识别这些瓶颈,并提出了三种代数重构方法来优化性能: 1. 片上Haar分析: 直接在芯片上重新计算低成本的Haar分析蝶形运算。 2. 闭式综合: 将多级综合级联折叠为由输出坐标位索引的单次遍历。 3. 权重折叠: 将学习得到的每个通道的缩放比例直接整合到卷积权重中。
This paper introduces an I/O-aware model to identify these bottlenecks and proposes three algebraic reformulations to optimize performance: 1. On-chip Haar Analysis: Recomputing the inexpensive Haar analysis butterfly directly on-chip. 2. Closed-form Synthesis: Collapsing the multi-level synthesis cascade into a single pass indexed by output-coordinate bits. 3. Weight Folding: Integrating learned per-channel scales directly into the convolution weights.
与参考实现相比,这些优化带来了 2.55倍的HBM流量减少、4.35倍的训练加速 以及 50%的峰值内存使用量减少,有效移除了此前限制WTConv效率的系统级开销。
These optimizations result in a 2.55× reduction in HBM traffic, a 4.35× training speedup, and a 50% reduction in peak memory usage compared to the reference implementation, effectively removing the systems-level overhead that previously limited WTConv efficiency.
Technical Contributions
The I/O Bottleneck
WTConv的参考实现是“内存受限(memory-bound)”的,这意味着其性能受限于处理器与HBM之间的数据传输速度,而不是计算吞吐量。通过对I/O需求进行建模,作者发现多级分解和综合过程会产生过多的内存流量。
The I/O Bottleneck
The reference implementation of WTConv is "memory-bound," meaning its performance is limited by the speed of data transfer between the processor and HBM rather than computational throughput. By modeling the I/O requirements, the authors identified that the multi-level decomposition and synthesis processes create excessive memory traffic.
Algebraic Reformulations
所提出的方法将标准的WTConv流水线转换为融合的、I/O高效的操作: * 融合执行: 通过将分析、卷积和综合步骤结合起来,该模型最大限度地减少了将中间张量写回HBM的需求。 * 基于坐标的索引: 对综合级联进行了重构,允许根据坐标位直接计算输出值,从而消除了多级缓冲的需求。 * 权重集成: 将学习到的缩放因子折叠到卷积核中,减少了前向传播所需的运算次数和内存读取次数。
Algebraic Reformulations
The proposed approach transforms the standard WTConv pipeline into a fused, I/O-efficient operation: * Fused Execution: By combining the analysis, convolution, and synthesis steps, the model minimizes the need to write intermediate tensors back to HBM. * Coordinate-based Indexing: The synthesis cascade is reformulated to allow direct calculation of output values based on coordinate bits, eliminating the need for multi-stage buffering. * Weight Integration: Folding learned scaling factors into the convolution kernels reduces the number of operations and memory reads required during the forward pass.
Performance Evaluation
使用 WTConvNeXt 配置进行评估,重构后的方法展示出以下优势:
* 速度: 与参考实现相比,训练速度提升高达 4.35 倍。
* 内存: 峰值内存占用减少约 50%。
* 效率: 显著减少了HBM流量,使WTConv成为标准深度卷积(depthwise convolutions)的一种可行且高性能的即插即用替代方案。
Performance Evaluation
Evaluated using the
WTConvNeXtconfiguration, the reformulated approach demonstrates: * Speed: Up to 4.35× faster training compared to reference implementations. * Memory: Roughly 50% reduction in peak memory footprint. * Efficiency: Significant reduction in HBM traffic, making WTConv a viable, high-performance drop-in replacement for standard depthwise convolutions.
Access & Resources
Access & Resources
- View PDF
- HTML (Experimental)
- TeX Source
- License: CC BY 4.0
