Functional SAM

Implements Functional-SAM, a sharpness-aware update that perturbs only the network Jacobian and keeps the loss derivative at the unperturbed point.

For a loss \(L\) composed with a network function \(F\) (the concatenated logits), the SAM gradient at the perturbed point \(\theta+\rho\,\epsilon^*\) expands by the chain rule into a "logit path", which moves the loss derivative \(\nabla_F L\) along the perturbation, and a "functional path", which moves the Jacobian \(\nabla_\theta F\). The paper argues the logit path drives spurious sharpness minimization that fails to improve generalization on harder problems. Functional-SAM discards the logit-path contribution: it evaluates the network Jacobian at the perturbed parameters but multiplies it by the loss derivative taken at the original parameters, so only the functional path contributes to the descent direction.

\[ \begin{aligned} \epsilon^* &= \frac{\nabla_\theta L(\theta_t)}{\lVert \nabla_\theta L(\theta_t) \rVert} \\ g_t &= \nabla_\theta F(\theta_t + \rho\,\epsilon^*) \cdot \nabla_F L(\theta_t) \\ \theta_{t+1} &= \theta_t - \eta\, g_t \end{aligned} \]

where \(\theta\) are the parameters, \(\eta\) the learning rate, \(\rho\) the perturbation radius, \(\epsilon^*\) the normalized ascent direction, \(F(\theta)\) the network function (logits), \(\nabla_\theta F\) its Jacobian with respect to the parameters, and \(\nabla_F L\) the gradient of the loss with respect to the function outputs; standard SAM instead uses \(g_t = \nabla_\theta L(\theta_t + \rho\,\epsilon^*)\), which the chain rule splits into this functional path plus the discarded logit path.

Reference: Sidak Pal Singh, Hossein Mobahi, Atish Agarwala, Yann Dauphin, "Avoiding spurious sharpness minimization broadens applicability of SAM", arXiv 2025. https://arxiv.org/abs/2502.02407

---

Back to the Canon