[Paper Review] Rethinking Learning Rate Tuning in the Era of Large Language Models

Hongpeng Jin, Wenqi Wei+3 more|arXiv (Cornell University)|Sep 16, 2023

Topic ModelingComputer Science3 citations

TL;DR

This paper challenges conventional wisdom in learning rate (LR) tuning for Large Language Models (LLMs), demonstrating that standard policies like cosine decay often underperform compared to fixed learning rates. Using LRBench++, the authors benchmark LR policies across DNNs and LLMs, revealing that training loss is a poor proxy for LLM task performance and that fixed LR schedules can outperform adaptive ones, especially on ARC and HellaSwag benchmarks.

ABSTRACT

Large Language Models (LLMs) represent the recent success of deep learning in achieving remarkable human-like predictive performance. It has become a mainstream strategy to leverage fine-tuning to adapt LLMs for various real-world applications due to the prohibitive expenses associated with LLM training. The learning rate is one of the most important hyperparameters in LLM fine-tuning with direct impacts on both fine-tuning efficiency and fine-tuned LLM quality. Existing learning rate policies are primarily designed for training traditional deep neural networks (DNNs), which may not work well for LLM fine-tuning. We reassess the research challenges and opportunities of learning rate tuning in the coming era of Large Language Models. This paper makes three original contributions. First, we revisit existing learning rate policies to analyze the critical challenges of learning rate tuning in the era of LLMs. Second, we present LRBench++ to benchmark learning rate policies and facilitate learning rate tuning for both traditional DNNs and LLMs. Third, our experimental analysis with LRBench++ demonstrates the key differences between LLM fine-tuning and traditional DNN training and validates our analysis.

Motivation & Objective

• To re-evaluate the effectiveness of existing learning rate policies in the context of Large Language Model (LLM) fine-tuning, where traditional DNN assumptions may no longer hold.
• To identify the core challenges in learning rate tuning for LLMs, including the disconnect between training loss and task-specific performance metrics.
• To develop and validate LRBench++ as a benchmarking tool for systematic evaluation of learning rate policies across both traditional DNNs and LLMs.
• To provide empirical evidence that fixed learning rate schedules can outperform widely used adaptive schedules like cosine decay in LLM fine-tuning.
• To guide future research toward cost-effective, performance-driven LR tuning strategies tailored to the unique dynamics of LLM fine-tuning.

Proposed method

• Re-evaluate established learning rate policies (e.g., cosine decay, warmup, ReduceLROnPlateau) in the context of LLM fine-tuning, contrasting them with their use in traditional DNN training.
• Design and implement LRBench++, a benchmarking framework that standardizes hyperparameter configurations, training protocols, and evaluation metrics for both DNNs and LLMs.
• Integrate multiple evaluation metrics beyond training loss, including task-specific benchmarks like ARC and HellaSwag, to assess fine-tuned model performance.
• Conduct controlled experiments across multiple LLMs (e.g., Alpaca-7B, Vicuna-7B, WizardLM-7B) and datasets, comparing fixed, cosine, and warmup LR schedules.
• Use gradient-based optimization with Adam as the default optimizer, tracking parameter updates via the standard Adam update rule: $\Theta_{t+1} = \Theta_t - \eta(t) \frac{\hat{M}_t}{\sqrt{\hat{V}_t} + \epsilon}$.
• Analyze the correlation between training loss trends and downstream task performance, identifying non-monotonic, stair-step loss patterns that signal model improvement.

Experimental results

Research questions

• RQ1 How do standard learning rate policies designed for traditional DNNs perform when applied to LLM fine-tuning?
• RQ2 To what extent does training loss correlate with downstream task performance (e.g., ARC, HellaSwag) during LLM fine-tuning?
• RQ3 Can fixed learning rate schedules outperform adaptive schedules like cosine decay in LLM fine-tuning for specific benchmarks?
• RQ4 What are the key differences in optimization dynamics between LLM fine-tuning and traditional DNN training that affect learning rate policy effectiveness?
• RQ5 How can a standardized benchmarking framework like LRBench++ improve the systematic evaluation and selection of learning rate policies for LLMs?

Key findings

• Fixed learning rate schedules, such as a constant $2 \times 10^{-5}$, consistently outperform adaptive policies like cosine decay on the ARC and HellaSwag benchmarks, achieving the highest task-specific accuracy.
• Training loss during LLM fine-tuning exhibits a non-monotonic, stair-step decline pattern—particularly with significant drops at epochs 2 and 3—indicating that loss is not a reliable proxy for model performance.
• The training loss for a fixed LR of $2 \times 10^{-5}$ was higher than for other schedules, yet it achieved the best performance on ARC and HellaSwag, demonstrating a disconnect between loss and task accuracy.
• The widely adopted LR policy in practice (e.g., cosine decay with $2 \times 10^{-5}$ in Vicuna-7B and WizardLM-7B) does not yield optimal results, suggesting current defaults are suboptimal.
• State-based LR schedulers like ReduceLROnPlateau are ineffective for LLM fine-tuning because they rely on loss monitoring, which fails to reflect actual task performance.
• The experimental results confirm that LLM fine-tuning has unique optimization dynamics distinct from traditional DNN training, necessitating a rethinking of learning rate tuning strategies.