[论文解读] StrassenNets: Deep Learning with a Multiplication Budget

Michael Tschannen, Aran Khanna+1 more|arXiv (Cornell University)|Dec 11, 2017

Advanced Neural Network Applications参考文献 40被引用 11

一句话总结

该论文提出 StrassenNets，一种将深度神经网络中的矩阵乘法替换为基于和积网络（SPN）的可学习、低乘法次数算术电路的框架，实现乘法次数严格受限下的端到端训练。通过联合优化与知识蒸馏，该方法在保持全精度模型准确率的同时，实现了超过 99.5% 的乘法次数减少，并通过仅使用 7 次乘法而非 8 次，重新发现了 Strassen 的 2×2 矩阵乘法算法。

ABSTRACT

A large fraction of the arithmetic operations required to evaluate deep neural networks (DNNs) consists of matrix multiplications, in both convolution and fully connected layers. We perform end-to-end learning of low-cost approximations of matrix multiplications in DNN layers by casting matrix multiplications as 2-layer sum-product networks (SPNs) (arithmetic circuits) and learning their (ternary) edge weights from data. The SPNs disentangle multiplication and addition operations and enable us to impose a budget on the number of multiplication operations. Combining our method with knowledge distillation and applying it to image classification DNNs (trained on ImageNet) and language modeling DNNs (using LSTMs), we obtain a first-of-a-kind reduction in number of multiplications (over 99.5%) while maintaining the predictive performance of the full-precision models. Finally, we demonstrate that the proposed framework is able to rediscover Strassen's matrix multiplication algorithm, learning to multiply $2 imes 2$ matrices using only 7 multiplications instead of 8.

研究动机与目标

• 减少深度神经网络（DNNs）中的乘法次数，以实现资源受限硬件上的高效部署。
• 通过施加乘法预算，实现对算术复杂度与模型准确率之间权衡的细粒度控制。
• 从数据中端到端学习近似矩阵乘法变换，而非依赖于手工设计的算法（如 Strassen 或 Winograd）。
• 在乘法次数大幅减少的情况下，仍保持全精度模型的预测性能。
• 证明该框架能够重新发现已知的快速算法，如 Strassen 的矩阵乘法。

提出的方法

• 将 DNN 层中的矩阵乘法表示为两层和积网络（SPN），以解耦乘法与加法操作。
• 将 SPN 边缘权重表示为 {−1, 0, 1} 中的三元矩阵，通过非零权重的数量实现对乘法次数的显式控制。
• 使用数据驱动学习，通过端到端随机梯度下降联合优化 SPN 权重与网络参数。
• 应用知识蒸馏，将全精度教师模型的知识迁移至压缩的学生模型，实现最小的准确率损失。
• 通过将卷积和全连接层重新表述为基于 SPN 的变换，将该框架应用于卷积层和循环层，包括 ResNets 和基于 LSTM 的语言模型。
• 采用一种参数化策略，其中 SPN 中的隐单元数量（r）控制乘法预算，r 被设定为输出通道数或隐藏单元数。

实验结果

研究问题

• RQ1 可学习的、可微分的框架是否能比手工设计的算法（如 Strassen 或 Winograd）更有效地减少 DNN 中的乘法次数？
• RQ2 SPN 基变换的端到端学习在将乘法次数减少 99% 以上的同时，能在多大程度上保持模型准确率？
• RQ3 知识蒸馏是否能恢复乘法次数大幅减少的模型的预测性能？
• RQ4 该框架是否能重新发现已知的快速算法（如仅用 7 次乘法实现的 Strassen 2×2 矩阵乘法）？
• RQ5 与现有量化与压缩技术相比，该方法在乘法次数减少与准确率方面的表现如何？

主要发现

• 在 ImageNet 上的 ResNet-18 中，StrassenNets 实现了 99.63% 的乘法次数减少，同时与全精度模型相比仅造成 2.0% 的 top-1 准确率下降。
• 在语言建模任务中，该方法在 PTB 数据集上实现了 99.69% 的乘法次数减少，困惑度仅增加 3.3%。
• 通过知识蒸馏，该框架在图像分类和语言建模任务中均实现了超过 99.5% 的乘法次数减少，同时保持了全精度模型的预测性能。
• 在语言模型设置中，当 r ≥ n_out 时，蒸馏后的 StrassenNets 表现优于教师模型，表明知识蒸馏可增强压缩模型的性能。
• 该框架成功重新发现了 Strassen 的 2×2 矩阵乘法算法，学习到仅用 7 次乘法而非 8 次即可完成 2×2 矩阵乘法。
• 与具有相似乘法次数的 TWN 量化基线相比，该方法在困惑度上低 7.4%，证明了其在效率-准确率权衡上的优越性。