Neural Mastery
← Practice
Cosine Annealing Learning Ratemedium

Cosine Annealing Learning Rate

medium ML Systems & MLOps 15 min

NumPy allowed · plain Python earns 10 bonus points

Linear algebra step 1/4

A vector transforms

Matrix multiplication is the first computation behind every model.
Open lesson →
01vector
02matrix
03transform
04result
[2, −1] × [[1, 2], [−.5, 1]] = [0, −2]

Computed matrix × vector transformation

Implement cosine-annealing learning rate decay, one of the most widely-used real LR schedules for training large models.

Task

Implement cosine_annealing_lr(step, total_steps, base_lr, min_lr=0.0). Return min_lr + 0.5*(base_lr-min_lr)*(1 + cos(pi*step/total_steps)). Clamp step to total_steps if it exceeds it (so the schedule never goes below min_lr).

Function Signature

cosine_annealing_lr(step: int, total_steps: int, base_lr: float, min_lr: float = 0.0) -> float

Examples

Example 1: at step 0, full base_lr
Input: step = 0 total_steps = 100 base_lr = 0.1 min_lr = 0
Output: 0.1
Example 2: at final step, min_lr
Input: step = 100 total_steps = 100 base_lr = 0.1 min_lr = 0
Output: 0
Example 3: halfway point
Input: step = 50 total_steps = 100 base_lr = 0.1 min_lr = 0
Output: 0.05

Constraints

  • total_steps > 0
  • base_lr >= min_lr
Python3Saved ✓
import math

def cosine_annealing_lr(step, total_steps, base_lr, min_lr=0.0):
# Your implementation here
pass

Case 1: at step 0, full base_lr
Input: step = 0 total_steps = 100 base_lr = 0.1 min_lr = 0
Expected: 0.1
Case 2: at final step, min_lr
Input: step = 100 total_steps = 100 base_lr = 0.1 min_lr = 0
Expected: 0
Case 3: halfway point
Input: step = 50 total_steps = 100 base_lr = 0.1 min_lr = 0
Expected: 0.05