On-Device Fine-Tuning via Backprop-Free Zeroth-Order Optimization

TL;DR

This paper explores a memory-efficient zeroth-order optimization method for on-device fine-tuning of AI models, enabling larger models to operate within limited device memory by avoiding storage of intermediate states.

Contribution

It provides a theoretical analysis and empirical validation showing that MeZO allows larger models to be fine-tuned on-device compared to traditional backpropagation, given memory constraints.

Findings

MeZO enables larger models to fit in device memory during fine-tuning.

MeZO can achieve accuracy advantages over backpropagation under memory constraints.

Sufficient wall-clock time is necessary for MeZO to match or surpass backpropagation accuracy.

Abstract

On-device fine-tuning is a critical capability for edge AI systems, which must support adaptation to different agentic tasks under stringent memory constraints. Conventional backpropagation (BP)-based training requires storing layer activations and optimizer states, a demand that can be only partially alleviated through checkpointing. In edge deployments in which the model weights must reside entirely in device memory, this overhead severely limits the maximum model size that can be deployed. Memory-efficient zeroth-order optimization (MeZO) alleviates this bottleneck by estimating gradients using forward evaluations alone, eliminating the need for storing intermediate activations or optimizer states. This enables significantly larger models to fit within on-chip memory, albeit at the cost of potentially longer fine-tuning wall-clock time. This paper first provides a theoretical estimate of the relative model sizes that can be accommodated under BP and MeZO training. We then numerically validate the analysis, demonstrating that MeZO exhibits accuracy advantages under on-device memory constraints, provided sufficient wall-clock time is available for fine-tuning.