Learning Rate (In Depth)
Fully understanding the most important hyperparameter
The learning rate (LR) determines the step size in gradient descent. It's the most influential hyperparameter: no other parameter can affect training as strongly, positively or negatively.
Impact on the training run
⚡ Chaos
LR too large (e.g. 1e-2)
• Loss oscillates extremely
• No convergence
• Possible gradient explosion (NaN)
✓ Optimal
LR ideal (e.g. 2e-5)
• Loss drops steadily
• Stable convergence
• Reaches a low minimum
🐌 Too slow
LR too small (e.g. 1e-8)
• Barely any progress per epoch
• Very slow training
• Practically no learning
LR recommendations by model type and task
2e-5 to 5e-51e-5 to 3e-55e-5 to 1e-41e-4 to 3e-45e-5 to 1e-42e-4 to 3e-41e-5 to 5e-51e-3 to 3e-3The LR range test – finding it systematically
LR range test (implementation):
1. Initialize the model fresh
2. Train for 1 epoch with LR schedule: 1e-8 → 1e-1
(exponentially increasing over all steps)
3. Plot loss vs. LR (log scale)
4. Identify:
- Point where loss drops most steeply → LR_steep
- Point where loss starts to rise → LR_max
5. Good LR = LR_steep / 10 to LR_steep
Example result:
LR_steep = 5e-4
LR_max = 2e-3
→ Good LR = 5e-5 (safely below steep)
or 1e-4 (more aggressive)Linear scaling rule & batch-LR relationship
Linear scaling rule: LR_new = LR_base × (batch_size_new / batch_size_base) Example: Base: BS=8, LR=2e-5 New: BS=32 → LR = 2e-5 × (32/8) = 8e-5 For batch size < 512: linear scaling works well For batch size > 512: square root scaling is better: LR_new = LR_base × sqrt(BS_new / BS_base) Note: often optional for small changes (8→16).
Practical approach for FrameTrain
LR Scheduler Strategies
Intelligently steering the learning rate during training
An LR scheduler automatically adjusts the LR during training. Instead of a constant LR, it enables fast learning at the start and precise fine-tuning toward the end.
Learning rate scheduler: warmup + cosine decay
The 5 most important schedulers in detail
Warmup + Cosine Decay ⭐ (Gold standard)
⭐ RecommendedLR rises linearly from 0 to the target LR (warmup), then a gentle cosine curve down to nearly 0.
warmup_steps = 5–10% total; scheduler_type = "cosine"✓ Advantages:
Gentle cooldown at the end. Very stable, reproducible results. Widely used and well proven.
✗ Disadvantages:
Warmup steps must be defined in advance.
Warmup + Linear Decay
LR rises linearly (warmup), then a linear decline to 0.
warmup_ratio = 0.06; scheduler_type = "linear"✓ Advantages:
Simple, well documented, BERT standard.
✗ Disadvantages:
Cosine decay is marginally better for LLMs.
Warmup + Constant
LR rises linearly to the target LR, then stays constant.
warmup_steps = 100; scheduler_type = "constant_with_warmup"✓ Advantages:
Simplest configuration.
✗ Disadvantages:
No automatic cooldown – can become unstable toward the end.
Cosine Annealing with Warm Restarts (SGDR)
Repeats cosine decay cycles with restarts. After each cycle: new exploration.
T_0 = 100 (first cycle length), T_mult = 2 (extension)✓ Advantages:
Diverse checkpoints from different minima. Helpful for ensembles.
✗ Disadvantages:
More complex. The loss curve looks cyclical – can be confusing.
ReduceLROnPlateau
LR is automatically reduced when validation loss doesn't drop for N epochs.
factor = 0.5 (LR × 0.5 on plateau); patience = 2✓ Advantages:
Fully adaptive, no manual tuning.
✗ Disadvantages:
Can react too early or too late. Needs validation loss after every epoch.
Calculating warmup – a practical formula
Calculating warmup steps:
warmup_steps = total_steps × warmup_ratio (e.g. 0.06)
total_steps = (dataset_size / batch_size) × num_epochs
Examples:
1,000 examples, BS=8, 3 epochs:
total_steps = 125 × 3 = 375
warmup_steps = 375 × 0.06 = ~23 steps
5,000 examples, BS=16, 2 epochs:
total_steps = 313 × 2 = 625
warmup_steps = 625 × 0.06 = ~38 steps
Rule of thumb: 50–200 warmup steps for most jobsBatch Size & Gradient Accumulation
Balancing efficiency, memory, and quality
The batch size determines how many training examples are processed simultaneously before weights are updated. It directly affects memory requirements, training speed, and often generalization quality.
Small vs. large batch size
Small batch size (1–8)
- ✓ Low GPU memory usage
- ✓ Stochastic noise → sometimes better generalization
- ✓ More frequent weight updates
- ✗ Noisy gradients (fewer samples)
- ✗ Slower (less parallelism)
- ✗ Needs a lower LR
Large batch size (32–256)
- ✓ Faster training (better GPU utilization)
- ✓ More stable gradients
- ✓ Allows a higher LR (linear scaling)
- ✗ Much more VRAM
- ✗ Can generalize worse at very large BS
- ✗ Needs LR adjustment
Sweet spot for LLM fine-tuning
Gradient accumulation – the VRAM trick
Gradient accumulation resolves the dilemma between limited VRAM and a large effective batch size: you process small batches, accumulate gradients over several steps, and only then perform a weight update.
Gradient accumulation example:
Goal: effective batch size = 32, but only 6 GB VRAM
Batch size: 4 (fits in VRAM)
Gradient accumulation: 8 (accumulate over 8 steps)
Effective batch size: 4×8 = 32 ✓
Procedure:
Step 1: Forward(Batch1) → accumulate gradient ÷8
Step 2: Forward(Batch2) → accumulate gradient ÷8
...
Step 8: Forward(Batch8) → accumulate gradient ÷8
→ optimizer.step() (1 update from 32 samples)
→ optimizer.zero_grad()
Important: divide gradients by accumulation_steps!
Otherwise: updates become too largeVRAM recommendations by GPU
Calculating effective batch size
Effective batch size = batch_size × gradient_accumulation_steps
× num_gpus (for multi-GPU)
Example (1 GPU):
batch_size = 4, accumulation = 8, gpus = 1
→ effective = 4 × 8 = 32
Recommendation: aim for an effective BS of 16–32
If BS 32 doesn't fit in VRAM:
→ BS=4, accumulation=8 is equivalentOptimizer Comparison
Which optimizer for which situation?
The optimizer determines how gradients are turned into weight updates. Besides the LR, it's the most important factor affecting training speed, stability, and final performance.
The most important optimizers in detail
AdamW ⭐ (Standard)
lr=2e-5, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01Combines adaptive learning (Adam) with correctly decoupled weight decay. The difference from Adam: weight decay is applied directly to weights, not folded into the gradients. Result: better regularization.
✓ Advantages:
Stable, fast, robust. Works out of the box for almost all LLM fine-tuning jobs.
✗ Disadvantages:
Needs ~3× the memory of the model weights (moment statistics). For 7B = ~14 GB extra.
💡 ⭐ Always use AdamW as the first attempt
AdamW 8-bit (bitsandbytes)
bnb.optim.AdamW8bit(...)Quantizes the Adam moment statistics to 8-bit. Saves ~75% optimizer memory with minimal quality loss. Essential for full fine-tuning on 24GB GPUs.
✓ Advantages:
Cuts optimizer memory to a third. Quality nearly identical to fp32 AdamW.
✗ Disadvantages:
Requires the bitsandbytes library. Minimally slower due to dequantization.
💡 Use when GPU memory is tight for full fine-tuning
SGD with Momentum
lr=0.01–0.1, momentum=0.9, weight_decay=1e-4The simplest optimizer with momentum. Well understood and theoretically grounded. Often better for vision models (ResNet, etc.), but rarely better than AdamW for NLP/LLMs.
✓ Advantages:
Simple, less memory than Adam (no moment statistics). Theoretically well understood.
✗ Disadvantages:
Very sensitive to LR. Needs careful LR tuning. Almost always worse than AdamW for LLMs.
💡 For NLP/LLMs: always prefer AdamW
Adafactor
relative_step=True (auto-LR) or fixed LRSaves memory through factorized instead of full moment statistics. For very large models where AdamW needs too much memory. Standard for Google models (T5, PaLM).
✓ Advantages:
Dramatically less memory than AdamW. Scales to very large models.
✗ Disadvantages:
Requires more careful LR tuning. Convergence sometimes worse than AdamW.
💡 Fallback when AdamW causes OOM
Optimizer memory budget
Memory overview (for a 7B model, bf16):
Model weights (bf16): ~14 GB
Activations (BS=4): ~4 GB
Gradients: ~14 GB (same as weights)
Optimizer states:
AdamW (fp32): ~28 GB (2× weights for moments)
AdamW 8bit: ~7 GB (1/4 of AdamW fp32)
SGD: ~14 GB (only 1× moment)
Adafactor: ~2 GB (factorized)
Total (AdamW): 14+4+14+28 = ~60 GB → needs an A100!
Total (AdamW8b): 14+4+14+7 = ~39 GB → 2× RTX 3090 or A100
With LoRA (no gradient for frozen layers):
Total (AdamW LoRA): ~14+1+1+2 = ~18 GB → RTX 3090!Regularization
Techniques for better generalization and stable training
Regularization covers techniques that prevent the model from "memorizing" the training data, encouraging better generalization instead.
The most important regularization techniques
L2 Regularization (Weight Decay)
L_total = L_task + λ · Σ w²weight_decay = 0.01 (default), 0.1 (strong)Adds a term to the loss proportional to the square of the weights. Penalizes large weights and keeps the model "modest". Correctly decoupled in AdamW.
Effect: Prevents extreme weight values. Implicit "Occam's razor" – simpler solutions are favored.
💡 Increase weight_decay to 0.05–0.1 under strong overfitting.
Dropout
during training: h_i = 0 with probability pdropout = 0.1 (default), 0.3–0.5 (strong)Randomly deactivates neurons during training. The network can't rely on individual neurons and learns redundant representations.
Effect: The strongest regularization tool. Prevents co-adaptation of neurons.
💡 Only active during training! At inference: all neurons active (scaled).
Label Smoothing
y_smooth = y × (1-ε) + ε/Klabel_smoothing = 0.1 (default)Instead of hard 0/1 labels, soft targets are used (e.g. 0.9 for positive, 0.1/(K-1) for all others). Prevents an overconfident model.
Effect: Improves calibration (confidence = true probability). Slight overfitting reduction.
💡 Especially helpful for NLP tasks with a lot of label noise.
Gradient Clipping
if ||g|| > max_norm: g = g × (max_norm / ||g||)max_grad_norm = 1.0 (default)Limits the maximum gradient norm. Not really a regularizer, but stabilizes training and prevents explosions.
Effect: Prevents gradient explosion and NaN loss. Essential for large models.
💡 Standard in FrameTrain: always enabled. For very stable trainings: 0.5 for stricter clipping.
Early Stopping
Stop when val_loss >= best_val_loss for patience epochspatience = 3 (default), monitor = "val_loss"Ends training once validation loss stops improving. Implicitly prevents overfitting from training too long.
Effect: Prevents the overfitting phase through a timely stop. Loads the best checkpoint.
💡 Always combine with save_best_only=True for automatic best-checkpoint saving.
Combining regularization
Recommended combination for FrameTrain fine-tuning: # Standard (little overfitting): weight_decay = 0.01 dropout = 0.05 (LoRA dropout) early_stopping patience = 3 # Moderate (medium overfitting): weight_decay = 0.05 dropout = 0.1 label_smoothing = 0.05 early_stopping patience = 2 # Strong (heavy overfitting, little data): weight_decay = 0.1 dropout = 0.2 label_smoothing = 0.1 early_stopping patience = 2 lora_rank = 4 (fewer trainable parameters)