A Stable Whitening Optimizer for Efficient Neural Network Training

TL;DR

This paper introduces SPlus, a stable and efficient neural network optimizer building on Shampoo, which improves convergence stability, enables learning rate transfer, and accelerates training across diverse tasks.

Contribution

We propose SPlus, a novel optimizer that addresses stability and efficiency issues in Shampoo, with new techniques for bounded updates, shape-aware scaling, and iterate-averaging.

Findings

SPlus achieves comparable performance to Adam with fewer gradient steps.

It reduces training time by up to 38%.

It maintains stability across various training stages and tasks.

Abstract

In this work, we take an experimentally grounded look at neural network optimization. Building on the Shampoo family of algorithms, we identify and alleviate three key issues, resulting in the proposed SPlus method. First, we find that naive Shampoo is prone to divergence when matrix-inverses are cached for long periods. We introduce an alternate bounded update combining a historical eigenbasis with instantaneous normalization, resulting in across-the-board stability and significantly lower computational requirements. Second, we adapt a shape-aware scaling to enable learning rate transfer across network width. Third, we find that high learning rates result in large parameter noise, and propose a simple iterate-averaging scheme which unblocks faster learning. To properly confirm these findings, we introduce a pointed Transformer training benchmark, considering three objectives (language modelling, image classification, and diffusion modelling) across different stages of training. On average, SPlus is able to reach the validation performance of Adam within 44-58% of the gradient steps and 62-83% of the wallclock time.