[论文解读] Low-Rank Training of Deep Neural Networks for Emerging Memory Technology

Albert Gural, Phillip Nadeau+2 more|arXiv (Cornell University)|Sep 8, 2020

Advanced Memory and Neural Computing参考文献 38被引用 4

一句话总结

该论文提出低秩微调（LRT），一种计算高效、内存轻量的算法，将批量大小与辅助内存解耦，从而实现在非易失性内存（NVM）边缘设备上的低写入密度训练。LRT 通过使用低秩权重矩阵和梯度最大范数正则化，在将权重更新次数相比 SGD 最多减少 90% 的同时保持高精度，在适应性学习和迁移学习任务中，其在准确率和写入效率方面均优于标准方法。

ABSTRACT

The recent success of neural networks for solving difficult decision tasks has incentivized incorporating smart decision making "at the edge." However, this work has traditionally focused on neural network inference, rather than training, due to memory and compute limitations, especially in emerging non-volatile memory systems, where writes are energetically costly and reduce lifespan. Yet, the ability to train at the edge is becoming increasingly important as it enables real-time adaptability to device drift and environmental variation, user customization, and federated learning across devices. In this work, we address two key challenges for training on edge devices with non-volatile memory: low write density and low auxiliary memory. We present a low-rank training scheme that addresses these challenges while maintaining computational efficiency. We then demonstrate the technique on a representative convolutional neural network across several adaptation problems, where it out-performs standard SGD both in accuracy and in number of weight writes.

研究动机与目标

• 为解决新兴非易失性内存（NVM）系统在设备内训练中面临的高写入能耗和有限耐久性挑战。
• 减少训练过程中所需的权重更新次数，降低导致 NVM 寿命下降和能耗增加的写入操作。
• 实现极低辅助内存的训练，使资源受限的边缘设备（其片上内存有限）具备可行性。
• 在极端内存和写入约束下保持高模型准确率，尤其是在权重量化和低批量大小条件下。
• 在兼容边缘硬件的多样化适应性和迁移学习场景中，验证 LRT 的可行性。

提出的方法

• LRT 将全精度权重矩阵分解为低秩因子 L 和 R，降低参数量和写入操作次数。
• 采用基于低秩更新的在线随机梯度下降，仅更新 L 和 R 矩阵，而非完整权重矩阵。
• 引入梯度最大范数正则化以稳定训练并减少梯度噪声，尤其在低秩近似下表现更优。
• 采用流式批量归一化，以在在线、增量学习设置中维持训练稳定性。
• 应用条件数阈值（$\kappa_{th}$）跳过病态矩阵的更新，减少计算量且对准确率影响极小。
• 该算法支持有偏和无偏 LRT 变体，实证分析表明无偏 LRT 在全连接层上表现更优。

实验结果

研究问题

• RQ1 低秩权重分解能否减少在基于 NVM 的边缘设备上训练所需的权重更新次数？
• RQ2 在内存和能耗受限条件下，LRT 与标准 SGD 相比，在准确率和写入效率方面表现如何？
• RQ3 梯度最大范数正则化在高量化和低批量大小条件下的低秩训练中起到何种稳定作用？
• RQ4 在极端权重量化下，LRT 是否能维持性能，而 SGD 因梯度压缩而失效？
• RQ5 流式批量归一化是否能提升在线、增量学习场景中的收敛性和稳定性？

主要发现

• 在多个在线适应任务中，LRT 即使在权重更新减少 90% 的情况下，仍实现了比标准 SGD 更高的准确率。
• 在 ImageNet 上使用 ResNet-34 进行迁移学习时，无偏 LRT 将 top-1 准确率恢复至 72.1%，显著优于 SGD 和高方差基线方法。
• 仅更新偏置的 LRT 导致准确率下降 15–30%，凸显了在低秩分解中权重更新的关键作用。
• 使用流式批量归一化显著提升了训练稳定性，尤其在无归一化设置下，对收敛性有显著益处。
• 应用条件数阈值（$\kappa_{th}=100$）使计算量减少约 2 倍，且对准确率影响极小；更高的阈值（$\kappa_{th}=10^8$）未带来性能提升。
• LRT 在极端权重量化下表现出强鲁棒性，SGD 因梯度压缩而失效，而 LRT 通过高比特宽的 L 和 R 矩阵保留了微小的梯度信号。