🤖 machine learning

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

This paper demonstrates that memory-efficient zeroth-order optimization (MeZO) enables on-device fine-tuning of significantly larger models than conventional backpropagation by eliminating the need to store activations and optimizer states, thereby trading increased wall-clock time for superior accuracy under stringent memory constraints.

Original authors: Prabodh Katti, Houssem Sifaou, Sangwoo Park, Bipin Rajendran, Osvaldo Simeone

Published 2026-05-06

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Prabodh Katti, Houssem Sifaou, Sangwoo Park, Bipin Rajendran, Osvaldo Simeone

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a giant, incredibly smart library of knowledge (a Large Language Model) that you want to bring with you on a small, battery-powered device, like a smartphone or a smartwatch. You want this device to learn new tricks specifically for your needs right there on the spot, without needing to call a massive server in the cloud. This process is called on-device fine-tuning.

The problem is that the device's "brain" (memory) is tiny compared to the library.

The Old Way: The "Backpropagation" Backpack

The traditional method for teaching these models is called Backpropagation (BP). Think of this like a student trying to learn a new subject by taking a test, then immediately writing down every single thought, scratch paper, and intermediate calculation they made while solving the problem so they can review it later to see where they went wrong.

The Analogy: Imagine you are trying to solve a complex math problem on a tiny napkin. To use the old method, you have to keep a copy of every single step you took on a separate piece of paper for every single layer of the problem.
The Result: The "napkin" (your device's memory) fills up instantly. Because you have to save all those intermediate notes, you can only fit a very small, simple library on your device. If you try to bring a huge library, you run out of space before you even start learning.

The New Way: The "MeZO" Intuition

The paper introduces a new method called MeZO (Memory-Efficient Zeroth-Order Optimization). This method doesn't write down the intermediate steps. Instead, it uses a "trial and error" approach.

The Analogy: Imagine you are trying to find the best route through a maze. Instead of drawing a map of every turn you made (which takes up a lot of paper), you just take a step, see if you hit a wall, take a tiny step in a different direction, and see if that's better. You only remember the result of the step, not the entire thought process that got you there.
The Result: You don't need to carry a heavy backpack of notes. You can fit a much, much larger library on your device because you aren't wasting space on "scratch paper."

The Trade-off: Speed vs. Size

The paper makes a very specific claim about the trade-off:

Size Advantage: Because MeZO doesn't need to store all those intermediate notes, you can fit models that are at least 2 times larger than what the old method allows. If you have a long conversation (a long "context"), the advantage grows massively—up to 25 times larger.
- Simple Math: If the old method lets you fit a 3-billion-word dictionary, MeZO might let you fit a 13-billion-word dictionary on the same device.
Speed Cost: The downside is that MeZO is slower. Because it has to "guess and check" many times to figure out the right direction (instead of just looking at its notes), it takes longer to learn.
- The Paper's Finding: In their tests, the old method learned quickly but hit a "ceiling" because it was forced to use a small model. The new method (MeZO) learned slowly, but because it was using a much bigger, smarter model, it eventually ended up with better accuracy after a few hours.

What They Actually Tested

The researchers didn't just do math; they ran experiments on a powerful computer chip (an H100 GPU) to simulate a device:

They compared a small model trained with the old "heavy backpack" method against a much larger model trained with the new "lightweight" method.
The Outcome: Even though the new method took longer (about 2 hours to get good results), the larger model it trained was significantly smarter (82% accuracy) than the small model trained quickly (under 75% accuracy).
They also tested "sparse" training (only updating 1% of the model's brain). Even with this trick, the old method still needed more memory than the new method because the "scratch paper" (activations) was still the biggest problem.

The Bottom Line

The paper concludes that if you want to run a truly smart AI on a device with limited memory, MeZO is the better choice. It allows you to bring a "bigger brain" to the edge, provided you are willing to wait a little longer for it to learn. It turns the limitation of "small memory" into an advantage by changing how the learning happens, rather than just shrinking the model.

Technical Summary: On-Device Fine-Tuning via Backprop-Free Zeroth-Order Optimization

Problem Statement

On-device fine-tuning is essential for edge AI systems to adapt to specific agentic tasks and evolving user needs without relying on external servers. However, a critical bottleneck exists: conventional fine-tuning via backpropagation (BP) is significantly more memory-intensive than inference. BP requires storing intermediate layer activations and optimizer states for every layer, often demanding an order-of-magnitude more memory than a single forward pass. Consequently, edge devices with strict memory constraints are forced to deploy much smaller models for fine-tuning than they could for inference, severely limiting the capability of on-device AI.

Methodology

This paper proposes and analyzes Memory-efficient Zeroth-Order Optimization (MeZO) as a solution to the memory bottleneck. Unlike BP, MeZO estimates gradients using only forward evaluations (perturbations), eliminating the need to store intermediate activations or maintain optimizer states.

The authors' approach involves two main components:

Theoretical Memory Analysis: The paper derives mathematical estimates for the memory footprints of BP and MeZO on decoder-only Transformer models. The analysis accounts for:
- Parameter Count: Including multi-head attention (and Grouped Query Attention variants) and Feedforward Neural Network (FFN) layers.
- Activation Storage: Calculating the memory required to cache activations for BP, which scales with context length ( $N$ ), batch size ( $B$ ), and hidden dimension ( $D$ ).
- Optimizer States: Assuming a stateless SGD optimizer to isolate the impact of activation storage.
- Implementation Factors: Introducing a factor $L' \leq L$ to represent the maximum number of layers of activations stored by MeZO due to implementation-specific buffering, allowing for a realistic comparison.
Empirical Validation: The theoretical claims are validated through numerical experiments on a single H100 GPU using the BoolQ and MultiRC datasets. The study compares:
- Full Fine-Tuning: Comparing BP (using GPT2-medium) against MeZO (using Llama2-7B and Llama2-13B) under fixed memory budgets (~17 GB).
- Parameter-Efficient Fine-Tuning (PEFT): Evaluating sparse updates (fine-tuning only 1% of parameters) for both BP and MeZO using Qwen3 models to assess the impact of sparsity on memory and accuracy.

Key Contributions

The paper makes the following specific contributions:

Memory Analysis: A theoretical comparison demonstrating that, under fixed on-chip memory budgets, MeZO can accommodate models at least 2× larger than BP. This advantage scales with context length, reaching up to 25× for long contexts (e.g., $N=32768$ ) without checkpointing, and approximately 5× with activation checkpointing.
Scaling Insights: The analysis reveals that while parameter memory scales linearly with the number of layers ( $L$ ), activation memory in BP scales quadratically with context length ( $N$ ) and linearly with $L$ . MeZO avoids the quadratic context penalty, making it increasingly superior as context windows grow.
Empirical Validation: Numerical results confirm that while MeZO exhibits slower convergence in terms of wall-clock time, it achieves superior validation accuracy under strict memory constraints because it allows for the deployment of significantly larger, more capable models.
PEFT Evaluation: The study shows that even with sparse updates (1% parameter tuning), BP's memory savings are modest (~12–13%) because activation memory dominates total consumption. In contrast, MeZO continues to enable the fine-tuning of larger models (e.g., 4B vs. 0.6B) within the same memory budget, yielding better performance.

Results

Accuracy vs. Time Trade-off: Experiments on the BoolQ dataset show that MeZO converges slower than BP in wall-clock time. However, after approximately 2 hours of fine-tuning, MeZO (using a 13B parameter model) achieves ~82% accuracy, whereas BP (using a smaller 0.6B-1.5B equivalent model due to memory limits) remains below 75%.
Memory Efficiency: The ratio of memory requirements (BP/MeZO) increases with context length. For a context length of 32,768, MeZO requires roughly 25 times less memory than BP without checkpointing.
Sparsity Limitations: The paper demonstrates that reducing the number of updated parameters via PEFT does not significantly alleviate the memory burden of BP, as activation storage remains the dominant factor. MeZO, by avoiding activation storage entirely, maintains its advantage even when BP is made sparse.

Significance and Claims

The paper positions MeZO as a strong candidate for on-device fine-tuning scenarios, particularly for agentic edge systems and continual learning applications. The primary significance lies in shifting the constraint from "model size" to "fine-tuning time." By accepting a longer wall-clock training duration, edge devices can fine-tune models that are orders of magnitude larger than those possible with backpropagation, thereby unlocking higher performance capabilities directly on the device.

The authors conclude that MeZO effectively reduces training memory demands to the level of inference memory, enabling the deployment of larger models and longer context windows on resource-constrained hardware. They note that while the current analysis focuses on standard Transformer architectures, the conclusions can be generalized to modern variants (e.g., sliding-window attention, Mixture-of-Experts) and that further research is needed to specialize these findings for neuromorphic systems and dynamic sparsity.