arXiv cs.AI· Amit Aflalo, Shahaf E. Finder, Roy Amoyal, Eran Treister, Oren Freifeld·· 10 小时前AI 评分26
通过 I/O 感知重构实现快速且内存高效的小波卷积
Fast and Memory-Efficient Wavelet Convolutions via I/O-Aware Reformulation
AI 导读
研究团队为 WTConv 建立 I/O 模型定位 HBM 数据搬运瓶颈，并通过三项代数重构实现 I/O 感知的融合实现，将建模 HBM 流量降低约 2.55 倍。在 WTConvNeXt 配置上，该实现相比参考实现训练速度最高提升 4.35 倍，峰值内存占用约减半。
正文 · AI 翻译
摘要：小波卷积（WTConv）作为一种日益流行的标准卷积即插即用替代方案出现，它随着分解层数的增加以指数方式扩展网络的感受野，同时保持参数量线性增长。然而，其参考实现由于通过高带宽内存（HBM）进行过多数据移动而严重受限于内存带宽。我们开发了一个 WTConv 的 I/O 模型来刻画这一瓶颈，并用它指导三种代数重构：(1) 在片上重新计算开销很小的 Haar 分析蝶形运算，(2) 将多级合成级联折叠为按输出坐标位索引的单次闭式遍历，(3) 将学习到的逐通道缩放折叠进卷积权重。这些重构共同实现了一个 I/O 感知的融合实现，大幅减少了 HBM 流量。我们在不同分解层数和广泛的张量形状范围内评估了 WTConvNeXt 配置。尽管执行的算术量相当，参考 WTConv 却比它所替代的深度卷积慢得多。我们的重构将建模的 HBM 流量减少约 $2.55\times$，相比参考实现带来最高 $4.35\times$ 的训练加速，同时峰值内存使用量大约减半。因此，我们的重构保留了 WTConv 的优势，同时大幅减少了其执行时间和内存占用，消除了此前限制其实际效率的系统开销。
| 评论： | 已被 TMLR 接收 |
| 主题： | 计算机视觉与模式识别（cs.CV）；人工智能（cs.AI） |
| 引用为： | arXiv:2608.10805 [cs.CV] |
| （或 arXiv:2608.10805v3 [cs.CV] 用于此版本） | |
| https://doi.org/10.48550/arXiv.2608.10805 由 DataCite 颁发的 arXiv DOI |
提交历史
来自：Amit Aflalo [查看邮箱]
[v1]
2026年8月11日星期二 11:24:05 UTC（37 KB）
[v2]
2026年8月13日星期四 03:25:29 UTC（37 KB）
[v3]
2026年11月4日星期三 08:41:51 UTC（37 KB）
来源：arXiv cs.AI · arxiv.org