PROFIT: A Specialized Optimizer for Deep Fine Tuning

Thank you for your positive assessment of our work. We are very much aligned that the positives you pointed out (our convincing results and our method’s lightweightness) are principles at the heart of what we hoped to accomplish with PROFIT.

Including LWF as a baseline

We will plan to run an analysis against LWF for a revision, but for now can offer the following theoretical analysis: for each step of PROFIT, we update the model by two signals: the signal $-\Delta$ which is equivalent to the soft distillation target update within LWF, and an additional update $g\perp \Delta$ which is a regularized version of the standard new task update within LWF. Therefore, one can think of PROFIT as a flexible version of LWF where we constantly update the “frozen” model to the model state at $t-1$. This allows PROFIT to surpass a few notable limitations to LWF that limited its practical utility:

1. LWF requires a separate inference step on a model copy that must be kept around.
2. LWF is difficult to configure when a model has to be maintained through multiple fine-tuning iterations, as it is unclear what to use as the frozen “base” model through multiple steps. PROFIT always sets the frozen model to the model at timestep $t-1$, which removes ambiguity on frozen model choice.

In addition, since PROFIT is a drop-in optimizer, it is significantly easier to configure and run.

Ultimately, it is reasonable to think of PROFIT as a practical evolution of LWF which removes standard barriers and ambiguity, while preserving the core idea. Importantly, PROFIT accomplishes this while removing assumed access to material data (i.e. frozen model weights) that LWF relies on. As such, we originally felt the comparison would be not too instructive, similarly to how comparing LWF to fine-tuning a model with the original dataset intact would also be difficult to interpret. Comparison to standard optimizers which make the same set of assumptions on access to data is more apples-to-apples. However, we are aware that given how much our methodology is inspired by LWF, a comparison may be of interest and we are working hard on completing that comparison to include in the supplementary. We hope that in the meantime this additional context at least clarifies how PROFIT is theoretically positioned with respect to LWF and why we decided to omit the comparison in the main paper to keep the analysis clean.

Zero-shot accuracy on a standard LLM benchmark suite

We will plan to run experiments on MMLU for the camera-ready version. One benchmark we ran that provides insight into the amount of general information retained are the CIFAR baselines we present in Section 4.1. Although not an LLM baseline, these experiments show that PROFIT beats all other methods in terms of knowledge retention and performance on the transfer set. We note that PROFIT is not intended as a silver bullet against catastrophic forgetting, but is mainly designed to mitigate forgetting in such a way that optimizes learning on the transfer set. This design principle explains why we largely reported transfer set accuracy results. However, we agree that these results could provide additional context on another baseline for PROFIT and will begin to run those experiments.

Thank you for your positive reading and recognizing our work as “one of the first optimizers explicitly designed for fine-tuning pre-trained deep learning models”.

How PROFIT compares to standard proximal optimization techniques (e.g: L2 regularization around the pre-trained weight)?

An L2 regularizer around the pre-trained weights relies on a maximally uninformative Bayesian prior that will indiscriminately pull updates back towards the pre-trained point. However, L2 regularization does not take into account weight updates towards the fine-tuning task at all. We already know that such indiscriminate preference of the original convergence point is suboptimal, because in all fine-tuning/post-training settings the old convergence point is assumed to be a poor solution (otherwise we would not fine-tune at all). PROFIT resolves this issue by continuing to update the model “anchor” point after each step, which allows for dynamically useful updates throughout model training.

We also note that L2 regularization around the pre-trained weight would necessitate keeping a copy of the model weights around to be able to perform this regularization step. This both induces compute/memory overhead but also adds ambiguity when a model needs to go through multiple iterative stages of fine-tuning. L2 regularization also adds a magnitude hyperparameter while PROFIT handles update magnitude balancing automatically and with theoretical grounding. These restrictions make this type of L2 regularization scheme less practical. PROFIT does not require any additional data, which is why the comparison itself is less meaningful (just like it would be difficult to compare the L2 regularization algorithm to a situation where we have access to all the original dataset). We chose to primarily focus on comparisons between PROFIT and settings that assume access to the same cardinality of information to ensure comparisons are all apples-to-apples.

Thus, we do not believe L2 regularization is an effective proximal optimization technique but also is only possible in a privileged setting that assumes higher levels of information access than we have in the more general PROFIT setting. To the best of our knowledge, we are the first to examine the “proximal fine-tuning” setting with these data restrictions, which also is the most general setting for fine-tuning.

What is the actual impact of hyper-parameters on performance and forgetting?

We quantify the impact of hyper-parameters in the supplementary. In particular, we point the reviewer to Table 4 and 5 where we ablate the critical hyper-parameters of our method. We leverage these hyperparameters in all the discussed experimental settings.

Examples where proximal fine-tuning arise in practice

The proximal setting is very common in machine learning applications. In computer vision, this can be applied when the sensor modality stays constant through the lifecycle - expanding classification/segmentation taxonomies in self-driving, localization and action in new indoor environments for humanoid robots, new rooms/layouts/maps for AR/VR applications. In language models, this can be used to fine-tune a base model to support new tool-calls or function calls, or in post-training to ensure pre-trained knowledge isn’t forgotten while instruction-following capabilities are being introduced. In general, any setting where a model has to be updated due to collection of new data would also fall into the proximal setting. We note that even nonproximal fine-tuning settings can be made proximal with a mild warmup stage. In Section 4.2, we presented experiments on VTAB where naive application of PROFIT leads to poor results, but a simple warmup period (10 out of 100 epochs) transforms the VTAB fine-tuning setting into a proximal setting, and at that point PROFIT leads to significantly improved results. In other words, even non-proximal fine-tuning settings can be easily transformed into proximal settings, making PROFIT actually a fully general algorithm for fine-tuning.

Whether PROFIT could be extended where source and target distributions differ

We point the reviewer to Section 4.2, where we show our method can be mitigated in practice by using a short warm-up phase with standard fine-tuning before switching to PROFIT. Furthermore, we show this approach works well in Table 2.

Revise or elaborate on Figure 1

We discuss how PROFIT differs from standard optimizers in the caption of Figure 1. Standard optimizers which take successive steps through parameter space, while the reference optimizer (superscript “r” - for “reference” - in the figure) keeps track how the system tries to move away from good state ($x_0$). The good state is then restored ($x_0$) and then reference steps are translated to $x_0$. The main optimizer then takes a step and the gradient update is orthogonalized wrt to the translated reference vector. For further details, we point the reviewer to Alg 1 and Sec 3.1 where we elaborate the intuition and operational mechanism behind our method.

Thank you for your detailed and thorough rebuttal. I appreciate the time and effort you put into addressing all of my questions and concerns. However, I still have a few additional questions and points of clarification that I would like to raise.

• The claim that L2 regularization around the pre-trained weights is “maximally uninformative” is debatable, as it in fact corresponds to an informative Gaussian prior when centered on pre-trained parameters. While it indeed pulls updates indiscriminately toward the anchor point, this can be beneficial for stability and preventing overfitting. The assumption that the pre-trained point is always a poor solution is also questionable, as fine-tuning often serves to adapt an already strong model rather than replace it entirely.
• In addition, since PROFIT appears to operate on the basis of full fine-tuning, it would be important to clarify whether similar benefits could be achieved in combination with parameter-efficient fine-tuning (PEFT) methods.

Thank you for your detailed response to my questions. I appreciate the clarifications you provided. However, some of my concerns remain unaddressed. Since your proposed method is motivated from a proximal optimization perspective, I believe that a comparison against a standard L2-based proximal optimization baseline is particularly important. Accordingly, I will make a slight adjustment to my score while keeping it generally consistent.

We thank the reviewer for going through our response. We believe the above table, experiments in Sec 4.1, 4.2, 4.3, 4.4, and the supplementary showcase the effectiveness of our method as compared to fine-tuning with another method. We are happy to add this baseline to our camera-ready version if given the chance to proceed.

We are happy to engage in further discussion to clarify additional concerns.

Thank you for finding our paper “well written” and “supported by empirical results across a decent breadth of tasks”.

Was the choice of a 10-epoch AdamW warmup a result of a systematic hyperparameter search?

The VTAB strategy was a heuristic chosen to warmup for 10% of epochs and fine-tune for the remaining epochs. We tried some variants with slightly higher warm-up epochs and found similar performance. Therefore, when a practitioner is faced with finetuning on arbitrary distributions, we recommend starting with a warmup for 10-20% of the total epochs.

Paper does not fully disclose the reason for including results from several state-of-the-art methods (e.g., VPT, NOAH)

This is an astute observation. We include these results for completeness and to show that PROFIT is competitive with these SOTA methods. However, we emphasize that these methods induce additional trainable parameters and the comparison is not apples-to-apples. We will add a column on Table 2 to indicate when a method induces additional trainable parameters, format the table to show clearer grouping between different methodology classes, and also include discussion of this point in the main text.

How the effort to implement PROFIT compares to setting up other advanced optimizers

We outline the algorithm of PROFIT in Alg 1. This logic can be encapsulated in an PyTorch optimizer, and our implementation is ~20 lines of code.

Paper should disclose the spatial (memory) overhead

We discuss the memory complexity for our method in Table 8 of the supplementary. Storing an additional copy of states for the reference optimizer does increase memory complexity, which is a limitation of our method as we point out in Sec 2.2.6 of the supplementary.

While PROFIT is implemented with two optimizers in our experiments for generality and flexibility, this is not a requirement. In practice, one can reuse a single optimizer instance (e.g., SGD, AdamW, etc.) for both the main and reference model, provided care is taken to isolate their parameter updates. This reduces the additional memory footprint to effectively zero.

Moreover, PROFIT is compatible with parameter-efficient fine-tuning techniques such as adapters (e.g., LoRA, VPT). In Table 3 of the supplementary, we show that combining PROFIT with these methods yields consistent performance improvements while keeping memory usage low.

Thanks the author for the response. Table 3 in supplementary materials is insightful. I would like the authors to further explain to the unclear points in their response: The authors first mentioned "Storing an additional copy of states for the reference optimizer does increase memory complexity, which is a limitation of our method as we point out in Sec 2.2.6 of the supplementary." Is this inevitable? Then, the authors mentioned "This reduces the additional memory footprint to effectively zero." Is this saying the current implementation is further costing more gpu memory? Please verify this.

Nevertheless, I would like to suggest the authors provide a typical memory / time usage profiling for the typical tasks including diffusion modeling, larger language model such as models from 1B to 7B, under the setting of parameter-efficient fine-tuning. This will better help to understand the wider applicability of the proposed method. Note that, this is not a suggestion asking to provide results showing improved performance, but asking the profiling results to reveal the potential gpu memory cost / computing time analysis.

We thank the reviewer for considering our response. In terms of memory benchmarking, we noted earlier that our current implementation of PROFIT induces additional memory footprint because we treat both main and reference optimizers completely separately in order to better explore the hyperparameter space. In theory, it is possible to zero out the memory footprint by implementing PROFIT logic in the same optimizer instance. We went back and implemented this version, where the optimizers do not make separate copies of the model weight parameters but share a copy, and we show the memory footprint results below:

As we see, the memory footprint is now effectively zero. At baseline, PROFIT induces zero additional cost at inference, and we demonstrated now that it also can have zero footprint at training. The main reason we did not do this during our experiments is that a separate custom implementation has to be done for every (O_ref, O_main) pair, and so it was an impediment for hyperparameter search. However, once you know the right optimizer types (and O_ref should generally be SGD), this implementation can be used to fully eliminate memory overhead. We hope this clarifies this particular point about why we felt it was justified to tolerate the additional memory used in our experiments, and alleviates the reviewer's concern.

Thank you for your thoughtful review and recognition that our work “uniquely resolves catastrophic forgetting” and “integrates as a plug-and-play optimizer”.

Two bolded entries per column in Table 2

We include results from adapter methods such as VPT, NOAH, etc for completeness and to show that PROFIT is competitive with these SOTA methods. However, we emphasize that these methods induce additional trainable parameters and the comparison is not apples-to-apples. We will add a column on Table 2 to indicate when a method induces additional trainable parameters, format the table to show clearer grouping between different methodology classes, and also include discussion of this point in the main text.

Suggest the authors to ​​expand experiments

We thank the reviewer for the suggestion. We believe the experiments reported in Sec 4.1, 4.2, 4.3, and 4.4 are comprehensive, and although we are time-constrained and cannot run the requested experiments in time for this response, we are confident they will show the same trend as the other experiments presented. We note that benchmarks added for VTAB are a wide variety of classification tasks (Sec 4.2), and we also demonstrated positive results on VLMs in Sec 4.3, which represents a bleeding-edge modern computer vision task. Our experiments on DriveLM also show improvements in a setting that is tangential to computer vision and very high-dimensional. However, we will plan to add additional experiments in time for the camera-ready.

Thank you for your review and finding our method “simple and effective” with “extensive experimental results”.

If proximal fine-tuning assumption could be a strong assumption

The proximal setting is very common in machine learning applications. In computer vision, this can be applied when the sensor modality stays constant through the lifecycle - expanding classification/segmentation taxonomies in self-driving, localization and action in new indoor environments for humanoid robots, new rooms/layouts/maps for AR/VR applications. In language models, this can be used to fine-tune a base model to support new tool-calls or function calls, or in post-training to ensure pre-trained knowledge isn’t forgotten while instruction-following capabilities are being introduced. In general, any setting where a model has to be updated due to collection of new data would also fall into the proximal setting. We note that even nonproximal fine-tuning settings can be made proximal with a mild warmup stage. In Section 4.2, we presented experiments on VTAB where naive application of PROFIT leads to poor results, but a simple warmup period (10 out of 100 epochs) transforms the VTAB fine-tuning setting into a proximal setting, and at that point PROFIT leads to significantly improved results. In other words, even non-proximal fine-tuning settings can be easily transformed into proximal settings, making PROFIT actually a fully general algorithm for fine-tuning.

How to measure if the proximal fine-tuning assumption is satisfied

In our setup, the pre-training and fine-tuning datasets are known to be sampled from the same underlying distribution. We also do show situations where proximality fails - for example in VTAB (Sec 4.2). However, we also noted within our VTAB experiments that a short warmup cycle (10 out of 100 epochs) moves the VTAB setting into the proximality range and at that point PROFIT provides significant training benefits. While we find measurement of the proximality condition to be complex and very dependent on dataset and model properties, we at least provide a method to induce proximality in any setting which should work for general use-cases.

How to choose the number of epochs for an arbitrary pair of pre-training and fine-tuning datasets

The VTAB strategy was a heuristic chosen to warmup for 10% of epochs and fine-tune for the remaining epochs. We tried some variants with slightly higher warm-up epochs and found similar performance. Therefore, when a practitioner is faced with finetuning on arbitrary distributions, we recommend starting with a warmup for 10-20% of the total epochs.

Why the standard SGD is recommended

The reference optimizer is used to perturb the model away from the local optima to estimate the local loss landscape. By construction, $\Delta$, the pseudo-gradient (L155), minimizes the loss on the old task. Using optimizers which rely on other statistics would mean that $\Delta$ does not faithfully represent the direction to minimize old task loss, since it may be influenced by information gathered at another point in the loss surface. This property only holds true for vanilla SGD, and hence the reference optimizer should be set to SGD.

In other words, PROFIT’s core algorithm is powered by simple orthogonality relationships between the main and reference gradient updates. If the reference update is calculated by any other optimizer than SGD, it is too complex for the simple orthogonality calculus to still hold true.

Experiments are done with only one random seed

Experiments were run multiple times and we report error bars here. These will be added in a revision if we are allowed to proceed. CIFAR experiments have a standard error <= 0.01, DriveLM experiments have standard error within 0.2 on accuracy, VTAB experiments shows the mean accuracy over 3 repeats (as discussed in the VTAB paper), Waymo Open Motion Dataset experiments have standard errors within 0.005m for 3s metrics and 0.015m for 8s metrics.

Formal proof for Theorem 3.1

At a convergence point $\mathbf{x}_0$, the first-order gradient vanishes and to leading order the loss surface will look like $L(\mathbf{x}) \approx \mathbf{x}_0 + 0.5(\mathbf{x} - \mathbf{x}_0)^t\mathbf{H}_0(\mathbf{x} - \mathbf{x}_0)$ for positive definite hessian $\mathbf{H}_0$. The gradient within this region is therefore $\nabla L(\mathbf{x}) \approx \mathbf{H}_0(\mathbf{x} - \mathbf{x}_0) \equiv \mathbf{H}_0\Delta$. Because $\mathbf{H}_0$ is positive definite given that $\mathbf{x}_0$ was a critical point, this means that $(-\Delta)^t\mathbf{H}_0\Delta < 0$, which means the $-\Delta$ direction is a valid gradient descent direction. We apologize for the confusion - by “$a$ points in the direction of $b$” in the part of the proof you referenced we intended the meaning to imply “pointing in the same direction as,” which implies $\langle a, b\rangle > 0$ and shows that $\Delta$ is a valid gradient descent direction. We hope this more formal proof clarifies any ambiguity and will include it in the final version of the paper.

Difference between Sec 4.1 and Sec 4.2

In Sec 4.1, we demonstrate results on image classification when the pre-training and fine-tuning datasets are sampled from the same distributions. In Sec 4.2, we illustrate a case where this assumption may not hold, and demonstrate a strategy to use PROFIT in such cases.