Does SGD really happen in tiny subspaces?

TL;DR

This paper investigates whether training neural networks within the dominant subspace of the loss Hessian is feasible, finding that such projections do not hinder training and that the observed alignment is likely spurious, challenging previous assumptions.

Contribution

The study demonstrates that projecting SGD updates onto the dominant subspace does not improve training, revealing that the alignment with this subspace is misleading and not essential for effective training.

Findings

Projected updates onto the dominant subspace do not decrease loss.

Removing the dominant subspace component does not impair training.

Alignment with the dominant subspace is likely spurious and not causally beneficial.

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. We observe similar behavior across practical setups, including the large learning rate regime (also known as Edge of Stability), Sharpness-Aware Minimization, momentum, and adaptive optimizers. We discuss the main causes and implications of this spurious alignment, shedding light on the dynamics of neural network training.