Fine-Tuning Language Models with Just Forward Passes

Summary

Fine-tuning with backpropogation becomes infeasible for very large language models because it uses too much memory. While zeroth-order optimization uses far less memory and could in principle fine-tune the model with just forward passes, past theory suggested that the learning rate must scale down with the number of parameters, making convergence prohibitively slow. However, this paper finds that zeroth-order optimization actually performs quite well and converges quickly even on very large language models. They provide theory to explain this fast convergence, where they show that under an assumption they call "low effective rank," the learning rate scales down with the rank rather than the number of parameters. They also provide a memory-efficient implementation of zeroth-order optimization that they call MeZO, along with memory efficient zeroth-order versions of SGD with momentum and Adam. In experiments, the method performs similarly to backpropogation with 1/12 the memory usage, while outperforming in-context learning and linear probing.

Strengths

(1) The paper is well-written. (2) The method is simple and easy to understand. (3) The theory provides useful insights into why zeroth-order optimization works for fine-tuning large pre-trained models. (4) The experimental results are strong, and the appendix contains thorough ablations. (5) The idea that zeroth-order optimization works well for fine-tuning LMs seems practically useful and addresses a pressing need in the community for memory-efficient methods.

Weaknesses

The paper seems strong overall, and I support its acceptance regardless of whether the suggested experiments below are run or not during the rebuttal period. (1) From what I understand, the paper does not verify the low effective rank assumption empirically, nor is it verified in the papers cited (which either study the effective rank / Hessian spectra in non-LLMs, or study the LLMs but not the effective rank and instead study the intrinsic dimensionality of fine-tuning). Therefore, to justify the assumption, it seems useful to study the Hessian spectra of the downstream fine-tuning loss in LLMs, at whatever size is feasible. (2) Related to (1), to verify the theory and confirm that the effective rank is indeed the quantity that determines convergence rates, it seems useful to run simulated experiments where one constructs a synthetic model + data and varies the effective rank, and examines whether the convergence rate or gradient norm scales with the effective rank as predicted in the theory.

Questions

(1) While MeZO is much more memory efficient, is it slower than backprop in terms of wall-clock time? (The appendix does state that FT used 1K steps while MeZO used 100K steps in the experiments. How much faster is each step of MeZO compared to each step of backprop?) (2) It's a bit surprising to me that MeZO, MeZO-prefix, and MeZO-LoRA optimize at similar speeds (Figure 5 in the appendix). Do the backprop versions also optimize at similar speeds? And do the three parameterizations actually have similar effective ranks empirically?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

n/a

Summary

This paper proposes a new zeroth order optimizer, MeZO, for LM training. This is proposed as an improvement to ZO-SGD. The advantage of this approach is a 12x reduction in the amount of memory required for training compared to backpropagation. This enables the training of much larger models. The effectiveness of MeZO is shown across a range of benchmarks and model sizes. The results compare favorably to linear probing and in context learning.

Strengths

The MeZO technique stands to unlock substantial capability for LM training. This enables training of much larger networks. The compatibility with LoRA, prefix tuning are important use cases for many LM users. There is an ability for optimizing non-differentiable objectives which is compelling, and could be expanded in the future. The empirical behavior are coupled with a section on theory which effectively describes the both the expected behavior, but elaborates on the expected slow convergence by expanding the theoretical analysis to address low effective rank networks.

Weaknesses

While the analysis refers to the convergence rate of MeZO, there is a very brief treatment of convergence behavior in the paper (Appendix E.2) It might be helpful for this to be expanded and possibly compared to backpropagation, especially in the context of the presentation of Section 4 Theory.

Questions

Appendix A (and Section 3) demonstrate that promoting is crucial to MeZO performance. Why is this? Much of the other behavior is supported by a theoretical treatment, but this observation stands out as relatively uninterrogated.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

yes though fairly lightly. The observation about prompts being critical for training may be a limitation for some (new) tasks or datasets.

Summary

The paper present a new optimiser MeZo based on stochastic approximation using gradient perturbation. This optimiser is very memory efficient as it only requires to perform 2 forward passes with different deltas/epsilons on the parameters and multiple gaussian samplings. These algorithms allows "finetunning" large language models to specific tasks very efficiently ( up to 30B on a single A100) yielding between x4 to x12 memory reductions. Since the proposed algorithm is an optimizer it can be combined with other standard techniques such as LORA or prefix tunning. All this is applied in finetunning setups similar to ICL.

Strengths

The algorithm is a new application of well known stochastic gradient approximation but many times forgotten due to their slowness. It is very surprising that this algorithm works, and the authors provide a theoretical justification of why this could be working in this case. They also acknowledge that despite of what it might sound this approach only works in prompting fine-tunning scenario [Appendix B.2]. The results are insight full , the baselines are fair and the theoretical analysis is correct to the best of my knowledge. This can settle as an alternative to in context learning with prompting, by fine-tuning with a limited set of examples (512). The proposed technique seems to work on the benchmarks used and seems to surpass other techniques such as Zero-shot, LP, ICL, and approaches very close performance to fine-tunning. This can be set a cheap alternative to ICL for some tasks.

Weaknesses

While there was a huge and titanic effort in the paper there were many questions that steam from the technique. The first area which has not been explored too much and given as true is the need of having a prompt to apply MeZO. This key ingredient is not well studied but rather given on some preliminary experiments, e.g. Table 5. Why is MeZO not working w/o prompts, even some very simple prompts ? The need of the prompt also raises the question of how this is related to ICL , as there are some works that suggest ICL maybe doing some alike to fine-tuning or back-propagation though the attention weights. Is this combination of prompting and MeZo that is guiding the forward propagations ? is there a mixed cooperation between the prompt and the stochastic technique ? How much of the prompt is needed ? Another question would be how the different techniques behave as a function of the number of examples k. It would have been nice to see a plot for some models at least between ICL , MeZo, and possible ft on the selected tasks. Why have authors stopped at 512 ? why only 16 or 512 ? there are some dataset that contain more training examples. Why didn't they compare fine-tuning and MeZo in other setups with larger examples ? There is the relationship between the task itself and the optimiser. It is not clear to me, in which tasks this will work properly. I suspect given the prompting above that this might only work on low-perplexity tasks or task in which prompting or ICL can generate good results and not in other more complex tasks. Clearly this is maybe too much to address in the paper, but all aspects above point toward the little understanding the reader is left with on under which conditions this technique can be applied. The future work seems to already assume the MeZO algorithm is working and proved, but there is just an hypotheses and a very low link between the experiment conditions and the theory. The link is stablished as "We attribute these phenomena to the Hessian of the loss exhibiting small local effective rank." . It would have been nice to strengthen this connection with some experiments or computation. Could this explain when or how this algorithm is applied ? if we remove the prompt does this increases the H effective rank ? would other tasks exhibit larger ranks ? how can we reduce it for each of the tasks ? Please, I would kindly ask the authors to read above questions and discussion as a signal of the interest the paper brought to the research field and not as criticism on their very interesting work.

Questions

While I have many question none of them affect the paper directly. It would be nice however if some of the weakness could be discussed by authors: * in which task will this technique work? do you have a characterisation ? have you tried other tasks where it failed ? * why the Hessian hypotheses ? is this based on some experimental or preliminar analysis ? * when ft surpasses this techniques? when we have to ft on hundreds of thousands of examples ? * why it only works with prompting ?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

While the author do not focus on the limitations, they are more or less clear per previous analysis. The focus of the paper is more on the direction of stressing the surprise of the technique with all known drawbacks is actually workig.

Summary

The paper proposes an enhanced memory efficient zero-order optimization method named MeZO. MeZO only requires the same memory as inference time and thus can enable model tuning for large LMs with limited memory budget. The authors demonstrate the efficacy of MeZO on multiple NLP benchmarks compared with linear probing, in context learning and fine-tuning in few/low shot learning regime. The authors also provided detailed theoretical proof for MeZO.

Strengths

0 - The authors targeted a not well understood domain (zero order optimization) and open up new opportunities for future works. In the era of LLMs, this method can enable many future work especially for those who don't have access to large-scale GPU clusters. 1 - Comprehensive experiments and ablation studies on the proposed method. 2 - Strong theoretical support on the proposed method. 3 - Good writing and flow which makes the paper easy to follow and understand. 4 - Well-articulated future work.

Weaknesses

No major weaknesses. I left some comments in the questions section and hope the authors can answer and address.

Questions

- Can authors also report practical training time compared with standard fine-tuning (e.g., in terms of # steps/ second)? - Maybe report average numbers as well in experiment results (e.g., Table 1) - Interested to see how the performance of MeZO compares with fine-tuning in the scenario of full training data (rather than k = 100, 500, etc).

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The authors have discussed limitations in the conclusion section and aim to explore them in future work.

Summary

This work introduced a memory-efficient zeroth-order optimizer that can fine-tune large language models with the same memory footprint as inference, using only forward passes and gradient estimates. Comprehensive experiments across model types, scales, and tasks, showing that MeZO outperforms zero-shot, in-context learning, and linear probing, and achieves comparable performance to fine-tuning with backpropagation, while reducing memory cost by up to 12 times. Non-differentiable objectives that MeZO can optimize, such as accuracy or F1 score, which are usually not amenable to backpropagation. Theoretical insights that explain why MeZO can optimize LMs with billions of parameters, despite classical zeroth-order analyses suggesting otherwise.

Strengths

1. It proposes a novel and memory-efficient method to fine-tune large language models without backpropagation, which can save up to 12x memory compared to standard methods. 2. It demonstrates that the proposed method can achieve comparable or superior performance to fine-tuning with backpropagation across various tasks, models, and tuning techniques. 3. It shows that the proposed method can optimize non-differentiable objectives, such as accuracy or F1 score, which are useful for many applications. 4. It provides theoretical insights on why the proposed method can overcome the classical limitations of zeroth-order optimization and leverage the benefits of pre-training and task prompts.

Weaknesses

1. While the experiments demonstrate good performance on the language understanding tasks, it is unknown whether the method is also applicable to the generation tasks. 2. It relies on the assumption of low effective rank of the Hessian matrix, which may not hold for all loss functions. It would be great to have a discussion about the scope of application for the proposed method.

Questions

How is the training stability of MeZO? Does it an advantage over the backpropagation-based methods, especially for the large models?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

1. It would be better to have the experiment results on the generation tasks, e.g. translation and summarization. 2. I suggest the authors to have a discussion about the scope of application for the proposed method.