Accelerating Neural Network Training Along Sharp and Flat Directions

TL;DR

The paper investigates how neural network training dynamics interact with Hessian curvature by separating updates into a Dominant subspace and its Bulk. It introduces Bulk-SGD, which updates in the orthogonal complement of the top-$k$ Hessian directions, and demonstrates through ablations and switch-point analyses that Bulk-SGD can accelerate training but suffers from instability at large learning rates. To balance speed and stability, the authors propose interpolated gradient methods that blend Bulk and Dominant components, revealing a trade-off where Bulk enhances descent speed while Dominant directions promote stability and generalization. A Hessian-decomposition analysis shows that curvature energy concentrates in the Dominant subspace and is largely captured by the Generalized Gauss–Newton term $H_o$, offering a principled path toward curvature-aware optimizer design with practical implications for faster, more stable training.

Abstract

Recent work has highlighted a surprising alignment between gradients and the top eigenspace of the Hessian -- termed the Dominant subspace -- during neural network training. Concurrently, there has been growing interest in the distinct roles of sharp and flat directions in the Hessian spectrum. In this work, we study Bulk-SGD, a variant of SGD that restricts updates to the orthogonal complement of the Dominant subspace. Through ablation studies, we characterize the stability properties of Bulk-SGD and identify critical hyperparameters that govern its behavior. We show that updates along the Bulk subspace, corresponding to flatter directions in the loss landscape, can accelerate convergence but may compromise stability. To balance these effects, we introduce interpolated gradient methods that unify SGD, Dom-SGD, and Bulk-SGD. Finally, we empirically connect this subspace decomposition to the Generalized Gauss-Newton and Functional Hessian terms, showing that curvature energy is largely concentrated in the Dominant subspace. Our findings suggest a principled approach to designing curvature-aware optimizers.

Figures (9)

• Figure 1: Combined results from switch point experiments. Training was performed with an MLP on the MNIST-5k dataset.
• Figure 2: Instability of Bulk-SGD is caused by the sharpness explosion. Training was performed on MNIST-5k dataset. In \ref{['subfig:sgd_top_evals']} and \ref{['subfig:bulk_sgd_top_evals']}, we plot the 21 leading eigenvectors for SGD and Bulk-SGD respectively. For Bulk-SGD training Dominant subspace of dimension $k = 20$ was used. The 21 eigenvector is marked with black color.
• Figure 3: Acceleration experiment. SGD and Bulk-SGD training loss on the MNIST-5k dataset. Each run is initialized after SGD 6,000 SGD steps.
• Figure 4: Training neural network using interpolated gradient steps. Heatmaps are formed by training with different parameters $\alpha_{\mathrm{Dom}}$ and $\beta_{\mathrm{Bulk}}$ on CIFAR10-5k dataset. In \ref{['subfig:interpolated_expoeriments_train_loss_log']} we plot the training loss and in \ref{['subfig:interpolated_expoeriments_test_acc']} we plot the test accuracy. Grey cells indicate training setups that resulted in the diverging loss. For more detailed version see \ref{['fig:interpolation_experiment_test_acc_full', 'fig:interpolation_experiment_training_loss_full']}.
• Figure 5: Hessian energy analysis. MSE Loss and Tanh activation. The rest of the parameters are the same as in Fig.