Narrowing the Focus: Learned Optimizers for Pretrained Models

“However, the field has many other famous approaches, such as Shampoo optimizers.”

We appreciate the reviewer's suggestion. We chose VeLO as it is the SoTA learned optimizer at this time. Shampoo is a traditional optimizer, but is worthwhile to consider as a baseline. Due to resource constraints we chose a limited number of baselines. We compared to Adam since it is a popular, robust method which tends to perform well with little hyperparameter tuning.

“How is the proposed approach determine if the learned optimizer is useful to help the model to be trained to converged rather than reaching certain accuracy?”

While model convergence is an important aspect for an optimizer, guaranteeing global convergence is a challenge for most optimization algorithms, including widely used methods like Adam (see e.g. https://arxiv.org/abs/1904.09237). Our method, like most other meta-optimization approaches, focuses on practical performance improvement rather than theoretical convergence guarantees.

However, because of the interpretability of our method, we can demonstrate certain aspects of the final performance that hints at convergence of our model (see Fig.9 from the appendix). First, close to the end of training, the learning rate drops significantly, which aligns with the principle of gradually reducing the learning rate as the model approaches a minimum. Second, our learned algorithm switches from Adam to SGD, which helps with the convergence. This transition leverages the strengths of each optimizer: Adam for initial exploration in early stages, and SGD for fine-tuning and potential convergence to sharper minima in later stages. This strategy is consistent with established practices in deep learning (see e.g. https://arxiv.org/abs/1705.08292).

These clarifications and additions will be incorporated into the final version of the manuscript.

“It seems to be odd to focus only on achieving accuracy. It would be better to see if the model is converged.”

Accuracy is the preferred metric for the datasets we’ve chosen, but we have included corresponding figures and tables using loss as the metric in the appendix. This allows for a comprehensive evaluation of performance beyond just accuracy.

“Can you give more back envelop estimation on the computational cost versus the benefits of the proposed approach bringing in in the actual useful applications?”

Table 1 provides the memory and compute overheads for our optimizer. While overall training a meta-optimizer is computationally more expensive, we argue that the benefits of our meta-training approach outweigh the costs in several practical scenarios, for example.

1. Fine-tuning large models often involves thousands of iterations with different datasets or tasks. Assuming a typical fine-tuning run requires $X$ compute hours and our meta-trained optimizer improves performance by $Y\%$, amortizing the meta-training cost ( $Z$ compute hours) becomes beneficial after approximately $Z / (X * Y/100)$ fine-tuning runs. Given the scale of the model deployments (e.g. for LLM), this threshold can be easily surpassed.

2. In domains where even small performance gains translate to substantial real-world impact (e.g., medical diagnosis, fraud detection), the investment in meta-training is justified. A $X\%$ improvement in accuracy, even at a higher initial computational cost, can lead to significant downstream benefits that outweigh the initial investment.

3. In cases where data or compute resources are limited, efficiently utilizing the available resources is critical. Our method, by achieving faster convergence and potentially better performance with fewer fine-tuning steps, directly addresses this constraint. For instance, if our method reduces the required training steps by $X\%$, this translates to a proportional reduction in time and cost, making fine-tuning feasible in scenarios otherwise impossible.

We appreciate the reviewer for the careful review and for pointing out many small errors and typos that we missed. We will certainly address and fix them as well as carefully re-review for any others missed.

“One of the problems mentioned in the intro is about the short horizon and how far we need to move from the pretrained model to train it.”

Thank you for raising this important point. We understand the apparent contradiction between highlighting the short horizon problem and then focusing on fine-tuning. Allow us to clarify.

While fine-tuning pretrained models does typically involve shorter training times compared to training from scratch, it does not inherently eliminate the short horizon bias. There are situations where we need to adapt to a slightly longer optimization trajectory than what the meta-optimizer was originally trained on. For example, we might encounter a new task or dataset that requires a few more fine-tuning steps to reach optimal performance. A meta-optimizer trained on very short horizons might fail to generalize effectively, even if the absolute fine-tuning time remains relatively short. In our work, we demonstrate that our proposed optimizer is robust to this problem. Specifically, we show that it can successfully generalize to longer fine-tuning horizons than those it was explicitly meta-trained on (e.g. see Figure 3 meta-generalization of L3RS curves from solid to dotted lines).

“I didn’t find a comparison to ATMO”. Is the ablation at L476 a representation of ATMO?

Thanks for raising this. The ablation at L476 and "Global" in Table 2 are our closest comparisons to ATMO, as they reduce our method to a weighted combination across all layers, similar to ATMO's linear combination of SGD and Adam.

However, even here our method is meta-learned, the linear weighting and learning rates are chosen dynamically, while ATMO uses a static hyperparameter. Thus, "Global" provides an upper bound on ATMO's performance. We'll clarify this in the manuscript.

“The experiments seem to focus on one structure and one variant of that structure”

We acknowledge that evaluating on a wider range of architectures is an important consideration. We initially focused on ResNet-34 due to its widespread use and established performance as a benchmark in the field. Furthermore, we don’t expect any intra-model generalization, since our current focus is fine-tuning from a fixed checkpoint and we expect some degree of model-specific adaptation is expected and, in some cases, even beneficial. We will be sure to highlight this scope and the potential for future architectural exploration in the revised manuscript.

“To note, I'm asking for a description for the reasoning of choosing those baselines.”

Our choices were guided by both standard practice and the need to isolate specific contributions. Adam serves as a robust and widely used optimizer benchmark. Cosine Schedule reflects the increasing popularity of cosine learning rate decay in fine-tuning for its smooth convergence. While other schedulers exist, the cosine schedule's effectiveness and prevalence made it a representative choice. Constant Learning Rate isolates the benefits of a schedule versus a single, fixed or learned global learning rate. 'Const head' adds a minimal learnable scaling in cases the budget is limited. We will clarify this reasoning in Section 5 of the revised manuscript.

“I’d expected to see more datasets with varying granularity”

Thank you for raising this important point about the diversity of datasets used in our evaluation. We understand the desire to see L3RS applied to a wider range of datasets with varying granularities. Our current focus, however, is on demonstrating the effectiveness of L3RS on established, large-scale image classification benchmarks (ImageNet and PLACES). This allows us to conduct a more in-depth analysis (like ablation in studies in Section 6) and demonstrate the effectiveness of L3RS within this specific context. We will clarify this scope and the motivation for our dataset selection in the revised manuscript.

Thank you for your detailed answers. I have a few questions and points to raise.

For example, we might encounter a new task or dataset that requires a few more fine-tuning steps to reach optimal performance.

As mentioned by the reviewer Gurv, in longer horizons, the proposed approach only outperforms Adam in the early stages with smaller training steps and on longer horizons, Adam works fine. So based on your answer to that comment and mine, would that mean that the specific targeted problem here would be very short horizon + few-shot or continual learning?

We initially focused on ResNet-34 due to its widespread use and established performance as a benchmark in the field

I would accept this answer in the light of the paper having deeper insights. However, the cited works that have been mentioned as being close to the submitted work seem to experiment with more than one structure. [a] has experimented with ResNet18, ResNet34, and BERT. [b] has experiments for GPT2, ResNet18, ResNet50. So I'd at least expect some potential results with one other structure on commonly used datasets but without any other structures, I wouldn't be comfortable changing my score.

I'm happy with the answers to the rest of my initial questions. Although, my comments about the writing of the paper have not been addressed. I expected the manuscript to be updated with better writing by now.

[a] Landro N, Gallo I, La Grassa R. Combining Optimization Methods Using an Adaptive Meta Optimizer. Algorithms. 2021; 14(6):186. https://doi.org/10.3390/a14060186

[b] Almeida, D., Winter, C., Tang, J., & Zaremba, W. (2021). A Generalizable Approach to Learning Optimizers. ArXiv, abs/2106.00958.

“the proposed learned optimizer only outperforms Adam in the early stages with smaller training steps” … “Given that the final performance of L3RS is comparable to Adam and considering the extra computational costs, what are the specific scenarios or benefits that justify the use of learned optimizers?”

While L3RS's final performance is comparable to Adam in long-horizon training, its key benefit lies in rapid adaptation. This makes L3RS particularly valuable in scenarios like few-shot learning, transfer learning, or when models require frequent and efficient fine-tuning on new data. The initial meta-training cost is amortized over numerous adaptation tasks, leading to overall computational savings in dynamic environments.

“Experiments on additional tasks and models are necessary.“

We agree that evaluating our method on a broader range of tasks and models would be beneficial. Our choice of ImageNet and PLACES was driven by their recognition as challenging, large-scale benchmarks that effectively highlight the strengths of our method in a complex visual domain. Due to resource constraints, we focused on these core vision tasks for this study, but we intend to explore these avenues in future work.

“A more in-depth discussion of advanced adaptive learning methods, such as Adabelief, is helpful.”

We appreciate the reviewer's suggestion to include a discussion of AdaBelief. Our choice of Adam was motivated by its status as a widely adopted and effective adaptive optimizer, providing a strong and relevant baseline for our method. While a comprehensive comparison with all advanced adaptive methods is computationally intensive, we agree that including AdaBelief is valuable. We will perform experiments with AdaBelief and provide a detailed analysis of the results in the appendix, including a discussion of its performance relative to Adam and our proposed method.

“L3RS is not generalized and needs to be trained independently for each model architecture”

We agree that exploring the broader generalization of L3RS across diverse architectures is a valuable direction for future work. However, it's important to note that our current focus is on the specific and challenging problem of fine-tuning from a fixed checkpoint. In this domain, some level of model-specific adaptation is often necessary due to the inherent differences between pre-trained models and downstream tasks. While L3RS is designed to be adaptable, our primary contribution lies in its effectiveness within this specific fine-tuning setting. We plan to investigate broader architecture generalization in future research.

“From ImageNet <-> PLACES experiments it seems a trained L3RS can generalize to other data distributions, and it would be more beneficial to include more benchmarks to showcase that.”

We understand that more evaluations are always preferred. We chose ImageNet and PLACES as the evaluation datasets due to their complexity and we believe that the current experiments reasonably show the usefulness of the method. The ideal usage of the method is to meta-train the optimizer on the data-distribution it will be used on. This generalization outside of the training distribution is not always expected with meta-learning methods and showcases the robustness of the method.

“Experiments are only conducted on ResNet34, while including evaluation on other architectures and more diverse benchmarks (e.g. VeLO) would make the results more convincing.”

We agree that evaluations of other architectures would be ideal, but have left that to future work. The VeLO paper did a good job of this, but VeLO is a general optimizer and needed to show meta-test time generalization to other architectures. The Velodrome dataset they created is also not applicable here because that dataset samples model architectures.

“Finetuning for only 1000 steps may not well represent the practical finetuning setting.“

We agree that the optimal number of fine-tuning steps can vary depending on the specific application. However, our experiments demonstrate that L3RS is highly effective even with relatively short fine-tuning durations, achieving strong results with up to 2500 steps. Moreover, our results show that L3RS achieves significant performance gains even within the 500-2500 step range, which is outside of the meta-training distribution, indicating its efficiency and effectiveness in practical fine-tuning scenarios.

“Lack of a deeper analysis on how SGD/Adam update directions are favored by different layers/learning stages.”

There are some patterns that we see with how the SGD/Adam directions are chosen throughout training. Figures 4, 8, and 9 show averages through training. In particular, Figure 8 shows the average parameter movement across all layers, reveals several distinct phases: an initial warm-up period with an increasing learning rate, a period of relatively constant learning rate while transitioning from ADAM to SGD, a phase of rapid learning rate decay, and a final convergence to the SGD direction over the last ∼10 steps. While the precise interpretation of these parameter dynamics is challenging, the L3RS parameters are considerably more interpretable than those of most black-box learned optimizers.

“For practical use of L3RS, there needs to be a L3RS training recipe generalized across different model architectures, or pretrained L3RS checkpoints tuned for popular architectures.“

Thank you for raising this important point about the practical application of L3RS. We agree that generalizability and ease of use are key for adoption. Regarding hyperparameter selection, we designed our meta-training procedure with broad applicability in mind, and we believe the chosen hyperparameters will be effective across a range of scenarios. To clarify this, we will expand the paper with a detailed discussion of the rationale behind our hyperparameter choices and provide guidance on potential adjustments for specific use cases.

We also strongly agree with the suggestion of providing pretrained L3RS checkpoints. Our vision for L3RS is indeed to offer a collection of specialized optimizers, tailored for different architectures and tasks, rather than a single monolithic model, like VeLO. We plan to release pretrained checkpoints for several popular architectures to facilitate practical use. Finally, we find the suggestion of investigating few-shot fine-tuning very compelling. This approach could enable rapid adaptation of L3RS to new models and data distributions, and we will explore this direction in future work.