LR Scheduler
Implement a warmup + cosine decay learning rate scheduler.
The schedule has two phases:
- Warmup (steps 0 to warmup_steps-1): Linear warmup from 0 to max_lr lr=max_lr×warmup_stepsstep
- Cosine decay (steps warmup_steps to total_steps-1): Decay from max_lr to min_lr lr=min_lr+0.5×(max_lr−min_lr)×(1+cos(π×total_steps−warmup_stepsstep−warmup_steps))
Input: warmup_steps total_steps max_lr min_lr Output: Learning rate for each step (0 to total_steps-1), one per line, rounded to 6 decimal places.
Example:
2 6 0.001 0.0001
0.000000 0.000500 0.001000 0.000775 0.000325 0.000100
- The warmup phase starts from step 0 to 1 (since warmup_steps = 2), with a linear increase from 0 to max_lr = 0.001. At step 0, the learning rate is 0.001×20=0.000000, and at step 1, it's 0.001×21=0.000500.
- From step 2 onwards, the cosine decay phase begins, with the learning rate calculated as 0.0001+0.5×(0.001−0.0001)×(1+cos(π×6−2step−2)).
- Applying this formula for steps 2 to 5 gives the remaining learning rates: at step 2, lr=0.001000 (since cos(0)=1), at step 3, lr=0.000775, at step 4, lr=0.000325, and at step 5, lr=0.000100.
- These calculated learning rates are then rounded to 6 decimal places and output, one per line, resulting in the given sample output.
Constraints:
- 0 < warmup_steps < total_steps
- min_lr < max_lr
- Round to 6 decimal places
More from LLM 2: Training & Alignment
Background Knowledge
The problem involves implementing a learning rate scheduler, which is a crucial component in machine learning and deep learning. A learning rate scheduler is used to adjust the learning rate during the training process of a model. The learning rate determines how quickly the model learns from the training data. A high learning rate can lead to fast convergence but may also cause the model to overshoot the optimal solution, while a low learning rate can lead to slow convergence.
The problem specifies a warmup and cosine decay schedule. The warmup phase is used to gradually increase the learning rate from 0 to a maximum value, allowing the model to start learning from the data without overshooting. The cosine decay phase is used to decay the learning rate from the maximum value to a minimum value, allowing the model to fine-tune its parameters and converge to the optimal solution. The cosine decay schedule is a type of annealing schedule, which is commonly used in optimization problems to gradually reduce the learning rate.
The mathematical formulas provided in the problem description are used to calculate the learning rate at each step. The warmup phase uses a linear interpolation formula to increase the learning rate from 0 to the maximum value, while the cosine decay phase uses a cosine function to decay the learning rate from the maximum value to the minimum value. Understanding these formulas and how they are used to calculate the learning rate is essential to solving the problem.
Algorithm/Approach
The general approach to solving this problem involves implementing a loop that iterates over each step from 0 to the total number of steps. At each step, the algorithm checks whether it is in the warmup phase or the cosine decay phase and calculates the learning rate accordingly. The algorithm uses the provided formulas to calculate the learning rate at each step and rounds the result to 6 decimal places.
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.