Unpacking the Implicit Norm Dynamics of Sharpness-Aware Minimization in Tensorized Models
This work addresses the optimization and generalization challenges in tensorized models for researchers and practitioners, offering a more efficient alternative to SAM, though it is incremental as it builds on existing SAM analysis.
The paper tackled the underexplored behavior of Sharpness-Aware Minimization (SAM) in tensorized models by analyzing its norm dynamics, showing that SAM's implicit control of Norm Deviation is governed by core norm and gradient covariance, and proposed Deviation-Aware Scaling (DAS) which achieved competitive or improved performance over SAM with reduced computational overhead in experiments across tasks like tensor completion and model compression.
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.