SAM Perturbation And Step

Problem

Perform one Sharpness-Aware Minimization (SAM) step.

Normalize the gradient: norm = sqrt(sum(g * g)), then perturb w_adv = w + rho * grad / (norm + eps) elementwise. Update the weights with the gradient evaluated at the perturbed point: w_new = w - lr * grad_adv. Return [w_adv, w_new] as two lists.

Hint

The perturbation uses a unit vector; the final step uses grad_adv, not grad.

Test cases

case 1

in: [[1,2],[0.5,-0.5],[0.1,0.1],0.1,0.05,1e-8]

expected: [[1.0353553385593275,1.9646446614406725],[0.99,1.99]]

case 2

in: [[0,0],[3,4],[-1,-1],0.01,0.1,1e-8]

expected: [[0.05999999988000001,0.07999999984],[0.01,0.01]]

case 3

in: [[-1,0.5,2],[1,2,-2],[0,0,0],0.2,0.01,1e-8]

expected: [[-0.9966666666777778,0.5066666666444445,1.9933333333555556],[-1,0.5,2]]

Starter codePython

def sam_update(w, grad, grad_adv, lr, rho, eps):
    # Returns [w_adv, w_new]
    # Your code here
    pass

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