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 (
|
(_, (inner_state, _)), grad = jax.value_and_grad( |
)
And SAM calculates gradients in this noised model again (
|
(_, (_, 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?
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
And SAM calculates gradients in this noised model again (
sam/sam_jax/training_utils/flax_training.py
Line 549 in dae9904
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?