Summary
**Context**. The focus of the present paper is on-device fine-tuning (gradient computation and weight update **starting from a pre-trained model**) under limited memory budget. One way to cut the memory cost of storing the computational graph for gradient computation by standard backprop is the Memory Efficient Zeroth Order (MeZO) optimizer [Malladi et al, 2023], whereby a directional gradient is computed via weight perturbation (a.k.a SPSA): computing the loss $L$ difference yielded by two forward passes with weights differing by $\epsilon u$ estimates $\nabla L \cdot u$. Since it is a purely forward procedure, it obviates the need to cache activations to execute a backward pass.
**Core contribution**. The present paper proposes a quantized version of MeZO where weight perturbation, gradient computation and weight update are carried out on quantized quantities. The proposed algorithm, coined QZO-FF (Alg. 1), is tested against a variety of fine-tuning tasks (few-shot learning, cross-domain and in-domain adaptation), modalities (image and audio data), and architectures (convolutional, attention-based, recurrent), with several variants being explored (with fp8 / fp32 activations) and benchmarked against standard backprop. The efficiency of QZO-FF, both in terms of resulting performance and memory usage, is demonstrated.
**Paper outline**. More precisely:
- Section 2 provides background knowledge on memory-efficient backprop (2.1), forward-mode differentiation (2.2) and quantized training (2.3).
- Section 3.1 and 3.2 formalizes further "forward gradients" (3.1) and the SPSA / weight perturbation procedure to estimate them (3.2). A hardware-friendly extension of SPSA coined as "Sign-m-SPSA", which estimates $\text{sign}(\nabla L \cdot u) u$, , is introduced along with the resulting SGD update (3.2).
- Section 3.3 presents the core algorithmic contribution by combining SPSA / weight perturbation and weight quantization (Alg. 1). More precisely, weights and perturbation are statically, symmetrically quantized (e.g. their range are estimated and set once, before fine-tuning), with one scale for each ($\Delta_w$ and $\Delta_q$). Therefore: i) $\Delta_w$ and $\Delta_q$ are fixed, with weights and perturbations quantized with 16 and 8 bits respectively, ii) the integer part of the perturbed weights is accumulated in 32 bit, iii) the dequantized perturbed weights is quantized-dequantized back into 16 bits using the same $\Delta_w$ scale (Eq. 6). The Sign-m-SPSA gradient estimator is applied and quantized-dequantized using the perturbation scale ($\Delta_z$, Eq. 7) . Finally, the weight update itself is quantized, such that it happens in the quantized integer part and is rescaled by $\Delta_w$ (Eq. 8). Alg. 1 summarizes the procedure in the case where the number of perturbed directions at use is 1 ($m=1$).
- Section 3.4 presents several algorithmic "enhancements" of the QZO-FF algorithm to improve the optimization procedure itself or its memory footprint.
- Section 4 presents experimental results. First, few-shot learning is considered (4.1) on visual and audio data. Here, "FF" refers for short to "QZO-FF". A quantized version of FF, where 8 bits activations are used, is also tested. On vision, three architectures are tested (ResNet12, 18 and ViT tiny) on 5 different standard few-short learning datasets. Two scenarii are considered: full fine-tuning and linear probing. It is shown overall that FF always yields better performance than the zero-shot baseline and within 5%, accuracy-wise, to the BP baseline on 26/30 experiments, and that the ViT backbone yields the least degradation. On audio, a similar experiment is done with two architectures (CRNN, AST) on two audio datasets. On 11/16 experiments, FF accuracy is 5% off the BP baseline. Then, a cross-domain adaptation task (4.2) is considered, where the different algorithmic enhancements previously introduced (e.g. quantized FF, gradient averaging, "sharpness aware" scheme...) are tested. Most importantly, it is observed that quantizing weights to 8 bits jeopardize the FF algorithm. Finally, sector, 4.3 presents in-domain OOD adaptation using the same fine-tuning schemes (LP, D-VPT) with three levels of corruptions of the CIFAR-10 dataset as OOD datasets. In this setting, FF achieves comparable performance with BP.
Strengths
- The problem tackled is highly relevant to on-device training, pragmatic and builds upon recent work [Malladi et al, 2023].
- The proposed algorithm is sound and well-explained.
- There are a lot of experimental settings, data modalities and architectures being explored.
- The proposed technique is effective in providing a learning signal, effectively training models and yielding relatively good performance compared to the BP baseline.
Weaknesses
- It is unclear what is kept in full precision in the proposed procedure (see my questions below).
- On a related note, it is also unclear that the proposed algorithm enhancements don't offset the advantages of manipulating statically quantized quantities (see my questions below).
- The experiments aren't all sufficiently well explained, neither in the main nor in the appendix, which is frustrating because there is a lot of work done there and we fail to deeply understand the proposed setups. I would even say that there are almost too many different experimental setups. Under constrained time budget to write the paper, I would have prioritized a lesser number of better detailed experiments rather than a lot of them left unsufficiently explained.
- **There aren't any error bars in any table and figures**, although the authors ticked in their checklist that they reported error bars and provided appropriate information about the statistical significance of their experiments (L. 520). For lack of this, it is very hard to draw any clear conclusion in terms of comparison between the different algorithms at use, e.g. is there a statistically significant gap here, or are these two results within error bar? We don't know.
- I don't understand what the 2D plot of the loss landscape really brings here in terms of insights.
Questions
- L. 135: "in order to mitigate the noisy component of forward gradients estimated by SPSA, we propose sign-m-SPSA": do you have evidence that sign-m-SPSA results in less noisy gradients?
- Section 3.3: could you please clarify what is kept in full precision? I see at least three different quantities not being quantized: i) the scales $\Delta_w$ and $\Delta_z$, ii) the loss for each perturbed weights and therefore its difference, iii) the averaged gradient (Eq. 7). Most importantly, do you confirm that you need to accumulate gradients across each direction ($i=1 \cdots m$) in higher precision (32 bits I guess?) and then quantize-dequantize it using the scale of the perturbation $\Delta_z$? Your pseudo-algorithm only treats the case $m=1$ so it remains unclear how this all work when $m>1$ and you need to average gradients. **Could you please write a new pseudo-algorithm**, alike Alg.1, **in the case $m>1$**, highlighting with **two different color codes** the quantized (int8 and int16) and full precision (fp32) quantities?
- L.193-198 (momentum-guided sampling): I think that incorporating momentum into your approach is crucial. However, I don't understand neither how it works. What do you mean by "as training progresses, a history of the momentum $z$ is incorporated to guide the new sampling process"?
- L.199 (sharpness-aware perturbation): you mean an "extra step of **directional** gradient ascent"?
- L.204-210 (sparse updates): which sparsity scheme did you employ? top-k magnitude based scheme may be quite costly, if it boils down to ranking all the weights by their magnitude.
- L.211 (kernel-wise normalization): in this case we agree that $\hat{g}$ needs to be stored in full precision? Also, the computation of the norms of $z$ and $w$ is computationally expensive ($O(d)$, where $d$ denotes the dimension of $w$ or $z$), as expensive as it would be to dynamically recompute $\Delta_z$ and $\Delta_w$, which you avoided by statically quantizing them. Don't you lose the advantage here of using static scales if in any case you need to perform this $O(d)$ operations?
- L.216 (few-shot learning): "a few labeled samples are available", but how many? Could you please clarify the experimental setup?
- L.222: "16w8a" means 16 bits for weights and 8 for activations, correct? I may not take this for granted and clearly define this notation.
- Table 2: I would rather compute the **relative** accuracy degradation (acc_BP - acc_qFF / acc_BP) rather than the **absolute** accuracy degradation.
- Table 2: the accuracy degradation when employing FF in the FT setting compared to BP in the same setting is quite severe (11.08% gap), although you are using a relatively small architecture (ResNet12) with relatively small input dimensionality (32x32). Why is this the case?
- Table 2: **there are not any error bars**, which makes it hard to make any sense of a $\sim 0.2/0.5$ difference between two experiments.
- Could you please define precisely what you mean by "zero shot" (I assume no training at all?), "linear probing" (I assume only the last linear layer is learned?) and "full fine-tuning" (all parameters are learned)?
- L. 238 (audio benchmark): could you please detail the few-shot setup for this task, and the tasks themselves? It is important for people not familiar with this literature.
- L. 252: I really did not understand what "cross-domain adaptation" really is about. Could you please explain better what it is?
- L.256: what is "visual-prompt tuning with deep prompts"?
- Fig. 2: except for large discrepancies between bars, it is difficult to draw any conclusion from this figure **for lack of error bars**. Could you please add them?
- Which conclusions / insights do you really gain from plotting the 2D contours of the loss landscape in the different settings?
Limitations
See weaknesses and questions above.
If I had detailed description of each of the experimental setups tackled, a precise knowledge of what exactly is kept in full precision, how some of the algorithmic enhancements really work and error bars on all figures, I would be prone to increasing my score. I really want to encourage the authors to do so because I do believe that the core of this work is of interest.