Regularizing Deep Neural Networks with Stochastic Estimators of Hessian Trace

TL;DR

This paper introduces a new regularization technique for deep neural networks that penalizes the trace of the Hessian matrix, improving generalization and flat minima discovery by using an efficient stochastic estimator.

Contribution

The paper proposes a novel Hessian trace regularizer for deep neural networks, utilizing Hutchinson's estimator with dropout for efficient computation, outperforming existing regularizers.

Findings

Outperforms existing regularizers and data augmentation methods

Enhances generalization by promoting flat minima

Efficient Hessian trace estimation with dropout

Abstract

In this paper, we develop a novel regularization method for deep neural networks by penalizing the trace of Hessian. This regularizer is motivated by a recent guarantee bound of the generalization error. We explain its benefits in finding flat minima and avoiding Lyapunov stability in dynamical systems. We adopt the Hutchinson method as a classical unbiased estimator for the trace of a matrix and further accelerate its calculation using a dropout scheme. Experiments demonstrate that our method outperforms existing regularizers and data augmentation methods, such as Jacobian, Confidence Penalty, Label Smoothing, Cutout, and Mixup.