跳到正文
arXiv · 在线与高效推理· Jafar Badour; Maurice van Keulen; Elena Mocanu·· 7 天前精选AI 评分77

超越掩码稀疏:SNACK 在 GPU 上实现真正稀疏的神经网络

Beyond Masked Sparsity: SNACK Enables Truly Sparse Neural Networks on GPU

AI 导读

基于摘要分析。该研究针对稀疏神经网络与 Dynamic Sparse Training(DST)常以稠密张量叠加二值掩码实现、难以兑现实际资源节省的问题,提出仅存储和计算非零连接的 GPU 稀疏层 SNACK,并提供连接重构与稀疏反向传播的 PyTorch API。 核心机制是将稀疏性落实到存储与计算,而非只将权重置零。配套 SNACK-COO 是采用 COO 格式的自定义 SpMM CUDA 内核,其 batch 到 Streaming Multiprocessor 的映射针对小批量、高稀疏度场景优化。摘要未说明具体连接更新策略、损失函数或训练调度,不能据此判断其是否适用于所有 DST 方法。 作者报告,内核相对掩码稠密基线 Dense+Mask 最高加速 7 倍;在 95% 稀疏度下,与 cuSPARSE、Sputnik 和 FlashSparse 具有竞争力。在 90% 稀疏度的单层评估中,作者报告训练相对 Dense+Mask 和完全稠密层分别加速 8 倍、3.7 倍,推理分别加速 4 倍、2 倍,内存使用较稠密层减少 72%,并降低能耗,但摘要未给出能耗量化结果。端到端方面,作者报告 GPT-2 峰值训练内存最多减少 40%;在 99% 稀疏度下,graph-style inference 相对 Dense+Mask 加速 4.8 倍。上述内核、单层与端到端结果不能直接互换,具体 GPU、矩阵尺寸、批量、计时方式及模型质量保持情况尚未验证,实验协议、消融及统计信息也尚未验证。 平台推测(待验证):若 EEG 解码模型包含可高度稀疏化的线性层,SNACK 可能复用于小批量推理与训练资源优化;该假设依赖层尺寸、稀疏度及 GPU 执行条件,而非 EEG 信号本身。可复用稀疏存储与 SpMM 内核,需改造层替换和连接更新接口。EEG 低信噪比、跨被试变化及小数据可能使高稀疏度损害解码性能,小模型也可能被索引和内核启动开销抵消收益。可在固定 Cross-subject 划分、相同 GPU 和批量下,对比 SNACK、Dense+Mask 与稠密层的 Accuracy、延迟及峰值内存,并扫描稀疏度,检验资源收益是否伴随性能下降。已有 EEG 相关工作尚未检索确认。

阅读价值

值得核对 SNACK 如何把掩码稀疏转化为实际 GPU 计算与存储节省,尤其是其小批量、高稀疏度下的内核表现及单层收益向端到端模型收益的转化条件。

来源:arXiv · 在线与高效推理 · arxiv.org