Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improving Generalization of Deep Neural Networks by Optimum Shifting

Yuyan Zhou, Ye Li, Lei Feng, Sheng-Jun Huang

TL;DR

The paper presents Optimum Shifting (OS), a penalty-free framework that moves neural-network parameters from sharp to flat minima without increasing training loss by exploiting under-determined linear systems in linear layers. A theoretical bound shows the Hessian trace $\mathrm{tr}(\mathbf{H}_{L})$ is controlled by the last-layer weight norm $\|\boldsymbol{V}\|^2$, so OS seeks the minimum-norm solution within $\boldsymbol{A}\boldsymbol{V}=\boldsymbol{Z}$, yielding $\boldsymbol{V}^* = \boldsymbol{A}^T(\boldsymbol{A}\boldsymbol{A}^T)^{-1}\boldsymbol{Z}$. To reduce computational cost, the authors develop Stochastic Optimum Shifting (SOS) guided by Neural Collapse, enabling OS on small batches with the expectation that empirical loss remains stable on the full dataset. Empirical results across CIFAR-10/100, ImageNet, and VOC show reduced Hessian traces and improved generalization, including modest to meaningful accuracy gains and better detection metrics, with SOS being compatible with standard regularizers like weight decay and SAM. The work broadens the toolkit for improving generalization by focusing on constraint-based flatness shifts rather than additive penalties, and demonstrates practical effectiveness on image-based tasks, while outlining future directions for NLP tasks.

Abstract

Recent studies showed that the generalization of neural networks is correlated with the sharpness of the loss landscape, and flat minima suggests a better generalization ability than sharp minima. In this paper, we propose a novel method called \emph{optimum shifting}, which changes the parameters of a neural network from a sharp minimum to a flatter one while maintaining the same training loss value. Our method is based on the observation that when the input and output of a neural network are fixed, the matrix multiplications within the network can be treated as systems of under-determined linear equations, enabling adjustment of parameters in the solution space, which can be simply accomplished by solving a constrained optimization problem. Furthermore, we introduce a practical stochastic optimum shifting technique utilizing the Neural Collapse theory to reduce computational costs and provide more degrees of freedom for optimum shifting. Extensive experiments (including classification and detection) with various deep neural network architectures on benchmark datasets demonstrate the effectiveness of our method.

Improving Generalization of Deep Neural Networks by Optimum Shifting

TL;DR

The paper presents Optimum Shifting (OS), a penalty-free framework that moves neural-network parameters from sharp to flat minima without increasing training loss by exploiting under-determined linear systems in linear layers. A theoretical bound shows the Hessian trace is controlled by the last-layer weight norm , so OS seeks the minimum-norm solution within , yielding . To reduce computational cost, the authors develop Stochastic Optimum Shifting (SOS) guided by Neural Collapse, enabling OS on small batches with the expectation that empirical loss remains stable on the full dataset. Empirical results across CIFAR-10/100, ImageNet, and VOC show reduced Hessian traces and improved generalization, including modest to meaningful accuracy gains and better detection metrics, with SOS being compatible with standard regularizers like weight decay and SAM. The work broadens the toolkit for improving generalization by focusing on constraint-based flatness shifts rather than additive penalties, and demonstrates practical effectiveness on image-based tasks, while outlining future directions for NLP tasks.

Abstract

Recent studies showed that the generalization of neural networks is correlated with the sharpness of the loss landscape, and flat minima suggests a better generalization ability than sharp minima. In this paper, we propose a novel method called \emph{optimum shifting}, which changes the parameters of a neural network from a sharp minimum to a flatter one while maintaining the same training loss value. Our method is based on the observation that when the input and output of a neural network are fixed, the matrix multiplications within the network can be treated as systems of under-determined linear equations, enabling adjustment of parameters in the solution space, which can be simply accomplished by solving a constrained optimization problem. Furthermore, we introduce a practical stochastic optimum shifting technique utilizing the Neural Collapse theory to reduce computational costs and provide more degrees of freedom for optimum shifting. Extensive experiments (including classification and detection) with various deep neural network architectures on benchmark datasets demonstrate the effectiveness of our method.
Paper Structure (23 sections, 1 theorem, 48 equations, 10 figures, 5 tables, 1 algorithm)

This paper contains 23 sections, 1 theorem, 48 equations, 10 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related work
  3. Optimum Shifting
  4. Motivation
  5. Our Proposed Optimum Shifting Method
  6. Stochastic Optimum Shifting.
  7. Gaussian elimination.
  8. Experiments
  9. Applying SOS to Trained Models
  10. Applying SOS in Training Process
  11. Complexity Analysis
  12. Conclusion
  13. Appendix
  14. Proof of Theorem 3.2. with CNN
  15. Proof of Theorem 3.2. with ResNet
  16. ...and 8 more sections

Key Result

Theorem 3.2

For an $l$-layer MLP or convolutional (CNN, ResNet and Densenet) neural network, given the loss function $L$, the trace of Hessian can be bounded by: where $\bm{V}$ is the weight matrix of the last linear layer and $C_0, C_1, C_2$ are constants and independent of the last hidden layer's weight $\bm{V}$. Therefore, if $\Vert \bm{V}\Vert^2$ is minimized, both the upper bound and the lower bound of

Figures (10)

  • Figure 1: SOS algorithm during training
  • Figure 2: Top 100 eigenvalues of the Hessian matrix before and after SOS.
  • Figure 3: Applying SOS to pre-trained models on ImageNet dataset. We also use three PyTorch official pre-trained models: ResNet50, Shufflenet_V2, and ResNext50 to test the effectiveness of our method. We can see that both the top-1 and top-5 accuracy have been increased by SOS algorithm.
  • Figure 4: Exact loss difference before and after SOS during training.
  • Figure 5: Visualization of the Hessian trace for different models and dataset with and without SOS during training.
  • ...and 5 more figures

Theorems & Definitions (2)

  • Definition 3.1
  • Theorem 3.2