SGD with Partial Hessian for Deep Neural Networks Optimization

TL;DR

This paper introduces SGD with Partial Hessian (SGD-PH), a novel optimizer combining second-order channel-wise Hessian information with first-order SGD to improve deep neural network training stability and performance.

Contribution

The paper proposes a new compound optimizer, SGD-PH, that accurately extracts partial Hessian matrices for channel-wise parameters, enhancing optimization in deep neural networks.

Findings

SGD-PH outperforms traditional optimizers on image classification tasks.

Partial Hessian information improves convergence stability.

The method maintains good generalization performance.

Abstract

Due to the effectiveness of second-order algorithms in solving classical optimization problems, designing second-order optimizers to train deep neural networks (DNNs) has attracted much research interest in recent years. However, because of the very high dimension of intermediate features in DNNs, it is difficult to directly compute and store the Hessian matrix for network optimization. Most of the previous second-order methods approximate the Hessian information imprecisely, resulting in unstable performance. In this work, we propose a compound optimizer, which is a combination of a second-order optimizer with a precise partial Hessian matrix for updating channel-wise parameters and the first-order stochastic gradient descent (SGD) optimizer for updating the other parameters. We show that the associated Hessian matrices of channel-wise parameters are diagonal and can be extracted directly and precisely from Hessian-free methods. The proposed method, namely SGD with Partial Hessian (SGD-PH), inherits the advantages of both first-order and second-order optimizers. Compared with first-order optimizers, it adopts a certain amount of information from the Hessian matrix to assist optimization, while compared with the existing second-order optimizers, it keeps the good generalization performance of first-order optimizers. Experiments on image classification tasks demonstrate the effectiveness of our proposed optimizer SGD-PH. The code is publicly available at \url{https://github.com/myingysun/SGDPH}.