Noise Stability Optimization for Flat Minima with Optimal Convergence Rates

by   Haotian Ju, et al.

We consider finding flat, local minimizers by adding average weight perturbations. Given a nonconvex function f: ℝ^d →ℝ and a d-dimensional distribution 𝒫 which is symmetric at zero, we perturb the weight of f and define F(W) = 𝔼[f(W + U)], where U is a random sample from 𝒫. This injection induces regularization through the Hessian trace of f for small, isotropic Gaussian perturbations. Thus, the weight-perturbed function biases to minimizers with low Hessian trace. Several prior works have studied settings related to this weight-perturbed function by designing algorithms to improve generalization. Still, convergence rates are not known for finding minima under the average perturbations of the function F. This paper considers an SGD-like algorithm that injects random noise before computing gradients while leveraging the symmetry of 𝒫 to reduce variance. We then provide a rigorous analysis, showing matching upper and lower bounds of our algorithm for finding an approximate first-order stationary point of F when the gradient of f is Lipschitz-continuous. We empirically validate our algorithm for several image classification tasks with various architectures. Compared to sharpness-aware minimization, we note a 12.6 eigenvalue of the found minima, respectively, averaged over eight datasets. Ablation studies validate the benefit of the design of our algorithm.


page 1

page 2

page 3

page 4


How to escape sharp minima

Modern machine learning applications have seen a remarkable success of o...

S-SGD: Symmetrical Stochastic Gradient Descent with Weight Noise Injection for Reaching Flat Minima

The stochastic gradient descent (SGD) method is most widely used for dee...

Faster Perturbed Stochastic Gradient Methods for Finding Local Minima

Escaping from saddle points and finding local minima is a central proble...

Natasha 2: Faster Non-Convex Optimization Than SGD

We design a stochastic algorithm to train any smooth neural network to ε...

Surrogate Gap Minimization Improves Sharpness-Aware Training

The recently proposed Sharpness-Aware Minimization (SAM) improves genera...

Neon2: Finding Local Minima via First-Order Oracles

We propose a reduction for non-convex optimization that can (1) turn a s...

Restarts subject to approximate sharpness: A parameter-free and optimal scheme for first-order methods

Sharpness is an almost generic assumption in continuous optimization tha...

Please sign up or login with your details

Forgot password? Click here to reset