[논문 리뷰] Tensor-Compressed Back-Propagation-Free Training for (Physics-Informed) Neural Networks

Yequan Zhao, Xinling Yu외 4명|arXiv (Cornell University)|2023. 08. 18.

Model Reduction and Neural Networks인용 수 6

한 줄 요약

이 논문은 신경망 및 물리학 기반 신경망(PINNs)을 위한 백프로파게이션 없이 학습할 수 있는 프레임워크를 제안한다. 텐서 압축된 제로차수(ZO) 최적화를 사용함으로써, 분산 감소와 하이브리드 ZO 기울기 추정 전략을 결합하여 메모리 효율적이고 확장 가능한 학습을 가능하게 하며, 자동 미분 없이도 엣지 디바이스에서 실행 가능하다. MNIST에서 최신 기술 수준에 근접한 정확도를 달성하고, 20차원 헬름홀츠-자코비-벨리만 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를 종단 간(end-to-end)으로 처음부터 학습할 수 있도록 하기 위해.
• 기존 ZO 최적화의 높은 분산과 고차원 설정에서의 낮은 확장성 문제를 텐서 압축과 하이브리드 기울기 추정 전략을 통해 해결하기 위해.
• PDE 손실 평가를 위한 기울기 계산을 백프로파게이션 기반에서 밀도 없는 격자 기반 스틸 추정기로 대체하여 PINNs에 대한 BP-free 학습을 확장하기 위해.

제안 방법

• 고차원 파라미터 기울기를 낮은 차원의 텐서 부분공간으로 투영함으로써 ZO 기울기 분산을 크게 감소시키는 텐서 압축 기반 분산 감소 기법을 도입한다.
• 무작위 편향과 유한차분 근사 값을 조합하여 전방 전파 횟수를 최소화하면서도 기울기 정확도를 유지하는 하이브리드 ZO 기울기 평가 방법을 제안한다.
• PDE 손실 함수 내에서 기울기를 계산하기 위해 희박 격자 기반 스틸 추정기를 개발하여, 백프로파게이션 없이도 PINN 학습을 가능하게 한다.
• 모든 학습 단계에서 전방 전파 전용 평가를 수행함으로써 중간 활성값 저장이나 체인 룰 기반 기울기 계산이 필요 없도록 한다.
• 전방 전파만을 사용하여 손실 평가 및 파라미터 갱신을 수행함으로써, 표준 피드포워드 네트워크와 PINNs 모두에 이 프레임워크를 적용한다.
• 이중 단계 학습 과정을 적용한다: 먼저 기울기 기반 ZO 기울기 추정(Zo-SignRGE)을 사용한 거친 학습을 수행하고, 이후 ZO-CGE를 사용한 정밀 조정을 수행하여 수렴 성능을 향상시킨다.

실험 결과

연구 질문

• RQ1 백프로파게이션 없이도 실제 크기의 신경망을 처음부터 효과적으로 스케일링할 수 있는가?
• RQ2 대규모 모델에 대해 BP-free 학습을 실현 가능하게 하기 위해 ZO 기울기 분산을 얼마나 효과적으로 줄일 수 있는가?
• RQ3 하이브리드 ZO 기울기 추정 전략을 통해 학습에 필요한 전방 전파 횟수를 줄일 수 있는가?
• RQ4 PDE의 도함수 계산을 밀도 없는 격자 추정기로 대체함으로써 백프로파게이션 없이도 PINNs를 학습시킬 수 있는가?
• RQ5 엣지 기반 작업에서 정확도와 수렴 속도 측면에서 BP-free 학습 성능이 일阶 미분 방법과 비교해 어떻게 되는가?

주요 결과

• 텐서 압축된 ZO 접근법은 스파arsity 기반 ZO 학습보다 MNIST에서 훨씬 낮은 검증 손실을 달성하며, FC-3 네트워크는 TT-3 파라미터화를 사용해 4.25e-05의 손실을 기록한다.
• MNIST 데이터셋에서 BP-free 학습은 3층 완전 연결 네트워크를 사용해 3.58e-04의 검증 손실을 기록하며, 자동 미분 기반 일阶 방법과 유사한 성능을 보인다.
• 이 방법은 20차원 헬름홀츠-자코비-벨리만 PDE를 정확도가 자동 미분 기반 일阶 학습과 유사한 수준으로 성공적으로 학습시켰다.
• 희박 격자 기반 스틸 추정기는 백프로파게이션 없이도 손실 함수 내에서 기울기를 정확하게 근사함으로써 BP-free PINN 학습을 가능하게 한다.
• 하이브리드 ZO 학습 전략은 전방 전파 횟수를 줄이면서도 기울기 품질을 유지하여 자원 제약이 있는 플랫폼에서 효율성을 향상시킨다.
• 정밀 조정 단계에서 더 높은 쿼리 비용이 발생하더라도, 전체 프레임워크는 엣지 디바이스에서 종단 간 BP-free 학습을 가능하게 하여 현장에서의 AI 학습을 위한 길을 열어 놓는다.