Inverse-Square-Root Learning Rate

Problem

Compute the inverse-square-root schedule used by Transformers:

lr = base_lr * min(step^-0.5, step * warmup_steps^-1.5)

Return 0.0 when step <= 0, and skip the warmup term when warmup_steps is not positive.

Hint

The two terms swap dominance exactly at step = warmup_steps.

case 1

in: [0.1,10,0]

expected: 0

case 2

in: [0.1,10,5]

expected: 0.015811388300841896

case 3

in: [0.1,10,100]

expected: 0.010000000000000002

case 4

in: [0.1,0,4]

expected: 0.05

Starter codePython

def inverse_sqrt_lr(base_lr, warmup_steps, step):
    # Your code here
    pass

4 test cases · run your code against them in the editor