Problems
Loading...
1 / 1

Warmup + Linear Decay LR Schedule

Optimization
Easy

Most modern training pipelines start with a warmup phase where the learning rate gradually increases from zero, followed by a decay phase where it gradually decreases. This prevents early instability from large updates while still allowing the optimizer to escape sharp minima later.

Given a base learning rate, warmup steps, total steps, and the current step, compute the learning rate at that step.

Schedule

Warmup phase (current_step < warmup_steps): the learning rate increases linearly from 0 to base_lr.

lr=base_lr×current_stepwarmup_stepslr = base\_lr \times \frac{current\_step}{warmup\_steps}

Decay phase (current_step >= warmup_steps): the learning rate decreases linearly from base_lr to 0.

lr=base_lr×total_stepscurrent_steptotal_stepswarmup_stepslr = base\_lr \times \frac{total\_steps - current\_step}{total\_steps - warmup\_steps}
Loading visualization...

Examples

Input:

base_lr = 0.1, warmup_steps = 10, total_steps = 100, current_step = 5

Output:

0.05

Step 5 is in the warmup phase. lr = 0.1 * (5 / 10) = 0.05.

Input:

base_lr = 0.1, warmup_steps = 10, total_steps = 100, current_step = 55

Output:

0.05

Step 55 is in the decay phase. lr = 0.1 * (100 - 55) / (100 - 10) = 0.1 * 45/90 = 0.05.

Hint 1

Check whether current_step is less than warmup_steps to determine which phase you are in.

Hint 2

Both formulas are simple linear interpolations. No special math functions are needed.

Requirements

  • During warmup, linearly increase lr from 0 to base_lr
  • During decay, linearly decrease lr from base_lr to 0
  • At current_step = warmup_steps, lr should equal base_lr

Constraints

  • base_lr > 0, warmup_steps >= 0, total_steps > warmup_steps
  • 0 <= current_step <= total_steps
  • Return a single float
  • Time limit: 300 ms
Try Similar Problems