提出轻量级设备端训练方法，用混合零阶与一阶优化提升模型精度和效率。
ElasticZO: A Memory-Efficient On-Device Learning with Combined Zeroth- and First-Order Optimization
- 混合零阶与一阶优化，仅对部分层用反向传播，降低内存开销。
- 相比纯零阶方法，准确率提升5.2%-9.5%，内存增加仅0.072%-1.7%。
- 首次实现全整数运算的零阶训练，适合资源受限的边缘设备部署。
零阶（ZO）优化因其仅需前向传播且内存接近推理水平，成为边缘设备上替代标准反向传播（BP）训练的有效方案。本文提出针对全精度与8位量化深度神经网络的设备端学习方法ElasticZO与ElasticZO-INT8。ElasticZO介于纯ZO与纯BP之间，采用对最后几层使用BP、其余层使用ZO的策略。ElasticZO-INT8首次实现仅整数运算的零阶训练，通过新方法从整数交叉熵损失值中计算量化后的零阶梯度。在分类数据集上的实验表明，ElasticZO有效缓解了纯零阶方法收敛慢的问题，显著缩小与BP训练的准确率差距。相比纯零阶，其准确率提升5.2%-9.5%，内存开销仅增加0.072%-1.7%，支持微调与全训练。ElasticZO-INT8进一步将内存使用和训练时间减少1.46-1.60倍与1.38-1.42倍，且不损失准确率。结果表明，该方法在准确率与训练成本间取得更好权衡，凸显了零阶优化在设备端学习中的潜力。
原文摘要 · Abstract (English)
Zeroth-order (ZO) optimization is being recognized as a simple yet powerful alternative to standard backpropagation (BP)-based training. Notably, ZO optimization allows for training with only forward passes and (almost) the same memory as inference, making it well-suited for edge devices with limited computing and memory resources. In this paper, we propose ZO-based on-device learning (ODL) methods for full-precision and 8-bit quantized deep neural networks (DNNs), namely ElasticZO and ElasticZO-INT8. ElasticZO lies in the middle between pure ZO- and pure BP-based approaches, and is based on the idea to employ BP for the last few layers and ZO for the remaining layers. ElasticZO-INT8 achieves integer arithmetic-only ZO-based training for the first time, by incorporating a novel method for computing quantized ZO gradients from integer cross-entropy loss values. Experimental results on the classification datasets show that ElasticZO effectively addresses the slow convergence of vanilla ZO and shrinks the accuracy gap to BP-based training. Compared to vanilla ZO, ElasticZO achieves 5.2-9.5% higher accuracy with only 0.072-1.7% memory overhead, and can handle fine-tuning tasks as well as full training. ElasticZO-INT8 further reduces the memory usage and training time by 1.46-1.60x and 1.38-1.42x without compromising the accuracy. These results demonstrate a better tradeoff between accuracy and training cost compared to pure ZO- and BP-based approaches, and also highlight the potential of ZO optimization in on-device learning.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。