THE LEARNING-RATE WARMUP

Difficulty: Medium Tags: Optimization, Learning Rate Schedules, Transformers

Story

A transformer training run uses linear warmup followed by inverse-square-root decay. A schedule implementation has an off-by-one bug at the warmup boundary. Calculate the learning rate for a specified step using the exact schedule below.

The Math

Let base learning rate eta, warmup steps W, and step s >= 1:

eta_s = eta * min(s / W, sqrt(W / s))

Assume W >= 1.

Input Format

eta (float), warmup_steps (int), step (int).

Output Format

Learning rate at that step (8 decimal places).

Constraints

• 0 < eta <= 1, 1 <= W, s <= 10^12.
• Time limit: 1.0 second.

Example

Input

eta = 0.001, warmup_steps = 4000, step = 1000

Output

0.00025000

Explanation: At step 1000, the warmup factor is 1000/4000=0.25, while the decay factor is sqrt(4000/1000)=2. The minimum is 0.25.