Stepping Forward on the Last Mile

TL;DR

This paper explores the use of fixed-point forward gradients for on-device training of deep learning models on resource-constrained edge devices, demonstrating its feasibility and proposing improvements to reduce memory and accuracy gaps.

Contribution

It investigates fixed-point forward gradient training for edge devices, introduces algorithm enhancements, and provides empirical analysis of its effectiveness across vision and audio tasks.

Findings

Feasible on-device training with fixed-point forward gradients demonstrated.

Proposed algorithm enhancements reduce memory footprint and improve accuracy.

Empirical results show competitive performance compared to traditional backpropagation.

Abstract

Continuously adapting pre-trained models to local data on resource constrained edge devices is the $\emph{last mile}$ for model deployment. However, as models increase in size and depth, backpropagation requires a large amount of memory, which becomes prohibitive for edge devices. In addition, most existing low power neural processing engines (e.g., NPUs, DSPs, MCUs, etc.) are designed as fixed-point inference accelerators, without training capabilities. Forward gradients, solely based on directional derivatives computed from two forward calls, have been recently used for model training, with substantial savings in computation and memory. However, the performance of quantized training with fixed-point forward gradients remains unclear. In this paper, we investigate the feasibility of on-device training using fixed-point forward gradients, by conducting comprehensive experiments across a variety of deep learning benchmark tasks in both vision and audio domains. We propose a series of algorithm enhancements that further reduce the memory footprint, and the accuracy gap compared to backpropagation. An empirical study on how training with forward gradients navigates in the loss landscape is further explored. Our results demonstrate that on the last mile of model customization on edge devices, training with fixed-point forward gradients is a feasible and practical approach.