Decoupled Decay With Step Schedule

Problem

Apply one decoupled weight decay step with a step-decayed coefficient.

Let wd_t = wd * decay^step, then return x - lr * grad_step - lr * wd_t * x. The optimizer step and the decay are applied separately, as in AdamW-style optimizers.

Hint

decay^step shrinks the decay coefficient over training.

Test cases

case 1

in: [2,0.5,0.1,0.01,0.5,0]

expected: 1.948

case 2

in: [2,0.5,0.1,0.01,0.5,2]

expected: 1.9495

case 3

in: [0,-1,0.2,0.05,0.9,3]

expected: 0.2

case 4

in: [-1,0,0.1,0.1,0,5]

expected: -1

Starter codePython

def decoupled_decay_scheduled(x, grad_step, lr, wd, decay, step):
    # Your code here
    pass

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