QAT Fine-Tune Step

Problem

Take one QAT fine-tune step with the straight-through estimator: w_new = w - lr*grad inside the clipping range |w| <= clip, and w_new = w outside it. Apply element-wise and return the new weight list.

The signature is qat_finetune_step(w, grad, lr, clip=1.0).

Hint

Weights outside the clip range receive no gradient under the clipped STE.

Test cases

case 1

in: [[1,-1],[0.1,0.2],0.5,1]

expected: [0.95,-1.1]

case 2

in: [[2],[1],1,1]

expected: [2]

case 3

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

expected: [-0.1,0.4]

case 4

in: [[1,2],[1,1],0.5,0]

expected: [1,2]

Starter codePython

def qat_finetune_step(w, grad, lr, clip=1.0):
    # Your code here
    pass

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