Optimizers and Learning Rate Schedulers for PEFT

Selecting the right optimizer and learning rate schedule is fundamental to successful model training, and this holds true for Parameter-Efficient Fine-Tuning (PEFT). While PEFT significantly reduces the number of trainable parameters compared to full fine-tuning, the dynamics of updating these low-rank matrices or adapter layers require careful consideration of optimization strategies to ensure stable convergence and optimal performance. Common optimizers and scheduling techniques used in PEFT workflows are presented, highlighting choices that balance performance with memory efficiency.

Choosing an Optimizer for PEFT

The optimizer's role is to update the trainable parameters (e.g., LoRA matrices $A$ and $B$, adapter weights) based on the computed gradients. While many optimizers exist, certain choices have become standard for PEFT due to their empirical performance and adaptations for memory efficiency.

AdamW: The Reliable Standard

AdamW remains a highly popular and effective optimizer for deep learning, including PEFT. It adapts the learning rate for each parameter individually, incorporating momentum (tracking past gradients) and scaling based on past squared gradients. Its distinguishing feature compared to the original Adam optimizer is its improved handling of weight decay. Instead of mixing weight decay with the gradient-based update, AdamW applies it directly to the weights after the main optimization step, which often leads to better generalization.

For PEFT, AdamW provides a baseline. Its adaptive nature helps navigate the sensitivities of different PEFT parameters. Typical configurations involve setting the learning rate, beta parameters ($\beta_1$, $\beta_2$), epsilon ($\epsilon$), and weight decay.

• Learning Rate: Often requires tuning. PEFT might tolerate slightly higher learning rates than full fine-tuning sometimes, but starting with values common for fine-tuning (e.g., $1e-4$, $5e-5$, $2e-5$) is a reasonable approach.
• Betas: $\beta_1$ (momentum decay) is usually set around 0.9, and $\beta_2$ (RMSprop decay) is often 0.999. These generally work well out-of-the-box.
• Epsilon: A small value ($1e-8$) for numerical stability.
• Weight Decay: A regularization term. Values like 0.01 or 0.1 are common starting points, but its impact might be less pronounced with PEFT due to the fewer trainable parameters. Tuning or even setting it to 0 is often viable.

Memory-Efficient Optimizers: 8-bit Adam

A significant advantage of PEFT is reducing memory requirements. Standard optimizers like AdamW maintain internal states (momentum and variance estimates) for each trainable parameter, typically in 32-bit floating-point precision. Even with fewer trainable parameters in PEFT, these optimizer states can still consume considerable GPU memory, especially when combined with gradient accumulation or large batch sizes.

To address this, memory-efficient optimizer variants have been developed. The most prominent example used with PEFT, particularly QLoRA, is 8-bit Adam. This variant quantizes the optimizer states, storing them using 8-bit data types instead of 32-bit. This drastically reduces the memory footprint of the optimizer, often by a factor of nearly 4x.

How it works:

1. The optimizer computes updates using 32-bit precision as usual.
2. The internal states (momentum, variance) are quantized to 8-bit representations for storage between steps. This involves techniques like block-wise quantization to maintain precision where it matters most.
3. Before the next update, the 8-bit states are de-quantized back to 32-bit precision for computation.

This approach provides substantial memory savings with minimal impact on convergence behavior or final model performance in most PEFT scenarios. Libraries like bitsandbytes provide implementations integrated with frameworks like PyTorch and Hugging Face's transformers. Using 8-bit Adam is often as simple as specifying a different optimizer name or class during trainer configuration.

Relative memory usage comparison for optimizer states. 8-bit Adam significantly reduces the memory required compared to standard 32-bit AdamW.

Other Optimizers

While AdamW and its 8-bit variant are prevalent, other optimizers like Adafactor can also be considered. Adafactor is designed for memory efficiency by factoring the second moment estimate matrix, avoiding the need to store the full matrix. It can offer memory savings without requiring explicit quantization libraries like bitsandbytes. However, its performance relative to AdamW can be more sensitive to hyperparameter choices, particularly the learning rate schedule.

Learning Rate Scheduling Strategies

The learning rate determines the step size taken during parameter updates. A fixed learning rate is rarely optimal. Learning rate schedulers dynamically adjust the learning rate during training, typically starting higher and decreasing over time. This helps achieve faster initial progress while allowing for finer adjustments as training converges.

Common LR Schedules for PEFT

1. Constant LR: The simplest approach, using a fixed learning rate throughout training. While easy to configure, it often leads to suboptimal results. It might be acceptable for very short fine-tuning runs or specific debugging scenarios.

2. Linear Decay: The learning rate decreases linearly from its initial value to a final value (often 0) over the course of training. This is a common and effective strategy.

3. Cosine Decay: The learning rate follows a cosine curve, decreasing from the initial value to a minimum value (e.g., 0). This schedule decreases the learning rate more slowly at the beginning and end of training and faster in the middle. It's often found to perform slightly better than linear decay in practice.

4. Warm-up Phase: It's highly recommended to combine decay schedules (Linear or Cosine) with an initial warm-up phase. During warm-up, the learning rate starts very low (often 0) and increases linearly to its target initial value over a specified number of steps (e.g., 5-10% of total training steps). This initial gentle phase helps stabilize training, preventing large, potentially disruptive updates early on when the model is rapidly adapting.

Learning rate progression for schedules with a 10-step warm-up phase followed by linear or cosine decay over 100 total steps. Initial learning rate after warm-up is normalized to 1.

PEFT-Specific Tuning Approaches

• Hyperparameter Sensitivity: The optimal initial learning rate, warm-up duration, and decay strategy can depend heavily on the specific PEFT method (LoRA vs. Adapters), the rank (r in LoRA), the dataset, and the base model. Experimentation is often necessary.
• Relationship with Batch Size: Effective batch size (global batch size = per-device batch size * gradient accumulation steps * number of devices) interacts with the learning rate. Larger effective batch sizes often allow for slightly higher stable learning rates.
• Training Duration: PEFT fine-tuning often requires fewer training steps compared to full fine-tuning. Ensure the LR schedule (especially decay and warm-up steps) is appropriately scaled to the total number of planned training steps.

Practical Implementation

Most deep learning frameworks and libraries like Hugging Face transformers provide easy ways to configure optimizers and LR schedulers. When using the Trainer API, you typically specify the optimizer type (e.g., 'adamw_torch', 'adamw_bnb_8bit'), learning rate, weight decay, and LR scheduler type (e.g., 'linear', 'cosine') along with warm-up steps within the TrainingArguments.

# Example using Hugging Face Trainer Arguments
from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4, # Effective batch size = 4 * 4 = 16
    optim="adamw_bnb_8bit",        # Use 8-bit AdamW
    learning_rate=2e-4,            # Initial learning rate
    lr_scheduler_type="cosine",    # Cosine decay schedule
    warmup_steps=50,               # Number of warm-up steps
    weight_decay=0.01,
    max_steps=500,                 # Total training steps
    fp16=True,                     # Use mixed precision (often used with PEFT)
    # ... other arguments
)

Choosing and tuning the optimizer and learning rate schedule is an iterative process. Starting with established defaults (AdamW or 8-bit Adam, cosine/linear decay with warm-up) provides a strong starting point. Monitor training loss and validation metrics closely to guide adjustments for achieving the best possible performance with your PEFT setup.