Implementation of the gradient approximation for the SAM objective (eq3)

[zhipeng-wei]

I have read your paper and code carefully.
You use the p as 2, thus SAM firstly maximizes the objective loss function by adding rho * normalized gradients to models' parameters (

sam/sam_jax/training_utils/flax_training.py

Line 537 in dae9904

(_, (inner_state, _)), grad = jax.value_and_grad(

)
And SAM calculates gradients in this noised model again (

sam/sam_jax/training_utils/flax_training.py

Line 549 in dae9904

(_, (_, logits)), grad = jax.value_and_grad(

).
In your paper, the gradient is calculated with respect to the original model's parameters, unlike the code which calculates gradient with respect to the noised model's parameters.

Are these two the same?