[论文解读] Tensor-Compressed Back-Propagation-Free Training for (Physics-Informed) Neural Networks

Yequan Zhao, Xinling Yu+4 more|arXiv (Cornell University)|Aug 18, 2023

Model Reduction and Neural Networks被引用 6

一句话总结

本文提出了一种无需反向传播的神经网络与物理信息神经网络（PINNs）训练框架，采用张量压缩的零阶（ZO）优化方法。通过结合张量压缩以降低方差，以及混合ZO梯度估计策略，该方法实现了在边缘设备上可扩展、内存高效的训练，无需自动微分，且在MNIST数据集上达到接近最先进（SOTA）的准确率，并成功求解了20维的哈密顿-雅可比-贝尔曼（Hamilton-Jacobi-Bellman）PDE。

ABSTRACT

Backward propagation (BP) is widely used to compute the gradients in neural network training. However, it is hard to implement BP on edge devices due to the lack of hardware and software resources to support automatic differentiation. This has tremendously increased the design complexity and time-to-market of on-device training accelerators. This paper presents a completely BP-free framework that only requires forward propagation to train realistic neural networks. Our technical contributions are three-fold. Firstly, we present a tensor-compressed variance reduction approach to greatly improve the scalability of zeroth-order (ZO) optimization, making it feasible to handle a network size that is beyond the capability of previous ZO approaches. Secondly, we present a hybrid gradient evaluation approach to improve the efficiency of ZO training. Finally, we extend our BP-free training framework to physics-informed neural networks (PINNs) by proposing a sparse-grid approach to estimate the derivatives in the loss function without using BP. Our BP-free training only loses little accuracy on the MNIST dataset compared with standard first-order training. We also demonstrate successful results in training a PINN for solving a 20-dim Hamiltonian-Jacobi-Bellman PDE. This memory-efficient and BP-free approach may serve as a foundation for the near-future on-device training on many resource-constraint platforms (e.g., FPGA, ASIC, micro-controllers, and photonic chips).

研究动机与目标

• 为了解决FPGA、ASIC和微控制器等边缘平台的硬件与内存限制，消除边缘设备上神经网络训练对反向传播的需求。
• 实现在无自动微分或梯度计算的前提下，从零开始端到端训练真实神经网络与PINNs。
• 通过张量压缩与混合梯度估计策略，克服传统ZO优化在高维设置下方差高、可扩展性差的问题。
• 通过将基于反向传播的导数计算替换为稀疏网格Stein估计器，将无反向传播训练方法扩展至PINNs。

提出的方法

• 提出一种张量压缩的方差减少技术，将高维参数梯度投影至低维张量子空间，显著降低ZO梯度的方差。
• 提出一种混合ZO梯度评估方法，结合随机扰动与有限差分近似，以最小化前向传播次数，同时保持梯度精度。
• 开发一种基于稀疏网格的Stein估计器，用于在PINN损失函数中计算导数，无需反向传播，从而实现无反向传播的PINN训练。
• 所有训练步骤均采用仅前向传播的评估方式，避免中间激活值存储或基于链式法则的梯度计算。
• 将该框架应用于标准前馈网络与PINNs，仅通过前向传播完成损失评估与参数更新。
• 采用两阶段训练流程：第一阶段使用基于符号的ZO梯度估计（Zo-SignRGE）进行粗调训练，第二阶段使用ZO-CGE进行微调，以提升收敛性。

实验结果

研究问题

• RQ1 零阶优化能否在无需反向传播的前提下，扩展至从零开始训练真实规模的神经网络？
• RQ2 如何充分降低ZO梯度方差，使无反向传播训练在大模型上成为可行？
• RQ3 混合ZO梯度估计策略能否减少训练所需的前向传播次数？
• RQ4 能否通过用稀疏网格估计器替代导数计算，实现无反向传播的物理信息神经网络训练？
• RQ5 在边缘设备相关任务中，无反向传播训练在准确率与收敛性方面与一阶方法相比表现如何？

主要发现

• 张量压缩的ZO方法在MNIST上的验证损失显著低于基于稀疏性的ZO训练，使用TT-3参数化时，FC-3网络的损失低至4.25e-05。
• 在MNIST数据集上，无反向传播训练使用三层全连接网络达到3.58e-04的验证损失，与使用自动微分的一阶方法相当。
• 该方法成功训练了PINN以求解20维哈密顿-雅可比-贝尔曼（Hamilton-Jacobi-Bellman）PDE，其精度接近使用自动微分的一阶训练结果。
• 稀疏网格Stein估计器通过在损失函数中准确近似导数，实现了无反向传播的PINN训练。
• 混合ZO训练策略在保持梯度质量的同时减少了前向传播次数，提升了资源受限平台的效率。
• 尽管微调阶段查询成本较高，整体框架仍实现了边缘设备上的端到端无反向传播训练，为设备端AI训练铺平了道路。