Unpacking the Implicit Norm Dynamics of Sharpness-Aware Minimization in Tensorized Models

TL;DR

This paper analyzes how Sharpness-Aware Minimization (SAM) implicitly regularizes tensorized models by controlling core norm imbalance and introduces a new method, Deviation-Aware Scaling (DAS), which improves performance and efficiency.

Contribution

The paper provides a theoretical analysis of SAM's implicit regularization in tensorized models and proposes DAS, a data-adaptive norm scaling method that enhances optimization.

Findings

DAS achieves competitive or better performance than SAM.

DAS reduces computational overhead compared to SAM.

SAM implicitly controls core norm imbalance through covariance of norms and gradients.

Abstract

Sharpness-Aware Minimization (SAM) has been proven to be an effective optimization technique for improving generalization in overparameterized models. While prior works have explored the implicit regularization of SAM in simple two-core scale-invariant settings, its behavior in more general tensorized or scale-invariant models remains underexplored. In this work, we leverage scale-invariance to analyze the norm dynamics of SAM in general tensorized models. We introduce the notion of \emph{Norm Deviation} as a global measure of core norm imbalance, and derive its evolution under SAM using gradient flow analysis. We show that SAM's implicit control of Norm Deviation is governed by the covariance between core norms and their gradient magnitudes. Motivated by these findings, we propose a simple yet effective method, \emph{Deviation-Aware Scaling (DAS)}, which explicitly mimics this regularization behavior by scaling core norms in a data-adaptive manner. Our experiments across tensor completion, noisy training, model compression, and parameter-efficient fine-tuning confirm that DAS achieves competitive or improved performance over SAM, while offering reduced computational overhead.