[论文解读] Low-rank lottery tickets: finding efficient low-rank neural networks via matrix differential equations

Steffen Schotthöfer, Emanuele Zangrando+3 more|arXiv (Cornell University)|May 26, 2022

Model Reduction and Neural Networks被引用 5

一句话总结

该论文提出了一种新颖的训练算法，通过使用矩阵微分方程在神经网络中动态保持低秩权重矩阵，实现了显著降低内存和计算成本的高效训练与推理。该方法在性能上可与全秩网络相媲美，同时在训练过程中自动适应秩的大小，展示了‘低秩彩票’的存在——即在训练过程中直接发现的高性能子网络，且不依赖于初始化。

ABSTRACT

Neural networks have achieved tremendous success in a large variety of applications. However, their memory footprint and computational demand can render them impractical in application settings with limited hardware or energy resources. In this work, we propose a novel algorithm to find efficient low-rank subnetworks. Remarkably, these subnetworks are determined and adapted already during the training phase and the overall time and memory resources required by both training and evaluating them are significantly reduced. The main idea is to restrict the weight matrices to a low-rank manifold and to update the low-rank factors rather than the full matrix during training. To derive training updates that are restricted to the prescribed manifold, we employ techniques from dynamic model order reduction for matrix differential equations. This allows us to provide approximation, stability, and descent guarantees. Moreover, our method automatically and dynamically adapts the ranks during training to achieve the desired approximation accuracy. The efficiency of the proposed method is demonstrated through a variety of numerical experiments on fully-connected and convolutional networks.

研究动机与目标

• 解决在资源受限环境中训练和部署大型神经网络所带来的高内存和计算成本问题。
• 克服现有低秩和剪枝网络方法的局限性，这些方法需要固定秩或依赖于训练后的压缩。
• 开发一种训练方法，能够在优化过程中自然发现并保持高效的低秩子网络（低秩彩票）。
• 确保在低秩流形上训练时的数值稳定性与收敛性保证，尤其是在小奇异值附近。
• 实现在训练过程中自动、动态地调整秩，消除对秩超参数的手动调优需求。

提出的方法

• 通过将权重矩阵表示为 $ W = U S V^ op $ 的形式，将权重矩阵限制在低秩流形上，其中 $ U $、$ S $ 和 $ V $ 为低秩因子，并在训练过程中进行更新。
• 将训练过程建模为在低秩流形上的连续时间梯度流，采用动力学低秩逼近（DLRA）技术求解矩阵ODE。
• 算法使用保持低秩结构的低秩数值积分器，确保在训练过程中即使在小奇异值附近也保持稳定性。
• 反向传播梯度通过矩阵链式法则计算，利用伴随法高效推导 $ U $、$ S $ 和 $ V $ 的更新规则。
• 对于卷积层，将核张量重塑为矩阵，并对展开后的核应用低秩分解，同时保持卷积结构。
• 该方法提供了在低秩流形上优化时的近似精度、下降行为和稳定性的理论保证。

实验结果

研究问题

• RQ1 是否能够直接在训练过程中高效发现并训练低秩子网络，而无需后期剪枝或依赖于初始化的搜索？
• RQ2 训练算法是否能在保持低秩结构的同时确保数值稳定性和收敛性，尤其是在奇异值较小时？
• RQ3 所提出的算法是否能在训练过程中自动适应权重矩阵的秩，以在最小化参数数量的同时达到期望的精度？
• RQ4 所得到的低秩子网络——即‘低秩彩票’——是否能在从随机初始化开始训练的情况下，实现与全秩模型相当的性能？
• RQ5 与标准的低秩训练和基于稀疏性的剪枝方法相比，该方法在效率和准确性方面表现如何？

主要发现

• 所提出的方法在全连接和卷积架构上实现了与全秩网络相当的测试准确率，即使参数量显著减少。
• 该方法在训练过程中动态调整权重矩阵的秩，消除了对人工设定秩超参数的需求。
• 数值实验表明，与全秩模型相比，该方法在训练和推理阶段均显著降低了内存使用量和计算成本。
• 该方法提供了在近似质量、下降行为和稳定性方面的理论保证，尤其在存在小奇异值时表现突出。
• 结果表明存在‘低秩彩票’——即高性能的低秩子网络——可直接在训练过程中被发现，且不依赖于初始化。
• 该方法优于基线低秩训练方案（如对 $ U $ 和 $ V $ 的交替SGD），通过在流形约束下优化，保持了更高的稳定性和准确性。