Does SGD really happen in tiny subspaces?

Minhak Song,Kwangjun Ahn,Chulhee Yun

2024-05-25

Abstract:Understanding the training dynamics of deep neural networks is challenging due to their high-dimensional nature and intricate loss landscapes. Recent studies have revealed that, along the training trajectory, the gradient approximately aligns with a low-rank top eigenspace of the training loss Hessian, referred to as the dominant subspace. Given this alignment, this paper explores whether neural networks can be trained within the dominant subspace, which, if feasible, could lead to more efficient training methods. Our primary observation is that when the SGD update is projected onto the dominant subspace, the training loss does not decrease further. This suggests that the observed alignment between the gradient and the dominant subspace is spurious. Surprisingly, projecting out the dominant subspace proves to be just as effective as the original update, despite removing the majority of the original update component. Similar observations are made for the large learning rate regime (also known as Edge of Stability) and Sharpness-Aware Minimization. We discuss the main causes and implications of this spurious alignment, shedding light on the intricate dynamics of neural network training.

Machine Learning,Optimization and Control

What problem does this paper attempt to address?

This paper investigates whether gradients in the training process of deep neural networks only update in the dominant subspace. The research found that although the gradients roughly align with the low-rank dominant subspace of the loss Hessian matrix during the training process, projecting the SGD updates onto this dominant subspace does not further reduce the training loss, indicating that this alignment may be misleading. Instead, projecting the updates onto the orthogonal complement of the dominant subspace, namely the "main quality subspace," is equally effective even though it removes most of the original updates. The paper points out that this misleading alignment is caused by the intrinsic stochastic noise in SGD and similar phenomena have been observed in other settings such as full-batch gradient descent and Sharpness-Aware Minimization. These findings reveal the complex dynamics in neural network training and imply the importance of the main quality subspace in the learning process.