A Stable Whitening Optimizer for Efficient Neural Network Training

• URL: http://arxiv.org/abs/2506.07254v2
• Date: Tue, 10 Jun 2025 22:01:14 GMT
• Title: A Stable Whitening Optimizer for Efficient Neural Network Training
• Authors: Kevin Frans, Sergey Levine, Pieter Abbeel,
• Abstract summary: Building on the Shampoo family of algorithms, we identify and alleviate three key issues, resulting in the proposed SPlus method.<n>First, we find that naive Shampoo is prone to divergence when matrix-inverses are cached for long periods.<n>Second, we adapt a shape-aware scaling to enable learning rate transfer across network width.<n>Third, we find that high learning rates result in large parameter noise, and propose a simple iterate-averaging scheme which unblocks faster learning.
• Score: 101.89246340672246
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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% of the gradient steps and 62% of the wallclock time.

Related papers

• Learning Hyperparameters via a Data-Emphasized Variational Objective [4.453137996095194]
  We consider direct gradient-based learning of regularization hyperparameters on the full training set via the evidence lower bound ("ELBo") objective from Bayesian variational methods.<n>In such scenarios, we find the ELBo prioritizes posteriors that match the prior variance, which leads to severely underfitting the data.<n>Our method reduces 88+ hour grid searches of past work to under 3 hours while delivering comparable accuracy.
  arXiv  Detail & Related papers  (2025-02-03T22:19:35Z)
• Why Line Search when you can Plane Search? SO-Friendly Neural Networks allow Per-Iteration Optimization of Learning and Momentum Rates for Every Layer [9.849498498869258]
  We introduce the class of SO-friendly neural networks, which include several models used in practice. Performing a precise line search to set the step size has the same cost during full-batch training as using a fixed learning. For the same cost a planesearch can be used to set both the learning and momentum rate on each step.
  arXiv  Detail & Related papers  (2024-06-25T22:06:40Z)
• Time Elastic Neural Networks [2.1756081703276]
  We introduce and detail an atypical neural network architecture, called time elastic neural network (teNN) The novelty compared to classical neural network architecture is that it explicitly incorporates time warping ability. We demonstrate that, during the training process, the teNN succeeds in reducing the number of neurons required within each cell.
  arXiv  Detail & Related papers  (2024-05-27T09:01:30Z)
• Class-Imbalanced Semi-Supervised Learning for Large-Scale Point Cloud Semantic Segmentation via Decoupling Optimization [64.36097398869774]
  Semi-supervised learning (SSL) has been an active research topic for large-scale 3D scene understanding. The existing SSL-based methods suffer from severe training bias due to class imbalance and long-tail distributions of the point cloud data. We introduce a new decoupling optimization framework, which disentangles feature representation learning and classifier in an alternative optimization manner to shift the bias decision boundary effectively.
  arXiv  Detail & Related papers  (2024-01-13T04:16:40Z)
• Adaptive Multi-step Refinement Network for Robust Point Cloud Registration [82.64560249066734]
  Point Cloud Registration estimates the relative rigid transformation between two point clouds of the same scene.<n>We propose an adaptive multi-step refinement network that refines the registration quality at each step by leveraging the information from the preceding step.<n>Our method achieves state-of-the-art performance on both the 3DMatch/3DLoMatch and KITTI benchmarks.
  arXiv  Detail & Related papers  (2023-12-05T18:59:41Z)
• Three Guidelines You Should Know for Universally Slimmable Self-Supervised Learning [4.631627683014556]
  We propose universally slimmable self-supervised learning (dubbed as US3L) to achieve better accuracy-efficiency trade-offs for deploying self-supervised models across different devices. We observe that direct adaptation of self-supervised learning to universally slimmable networks misbehaves as the training process frequently collapses. We propose three guidelines for the loss design to ensure this temporal consistency from a unified perspective.
  arXiv  Detail & Related papers  (2023-03-13T05:37:46Z)
• AdaSAM: Boosting Sharpness-Aware Minimization with Adaptive Learning Rate and Momentum for Training Deep Neural Networks [76.90477930208982]
  Sharpness aware (SAM) has been extensively explored as it can generalize better for training deep neural networks. Integrating SAM with adaptive learning perturbation and momentum acceleration, dubbed AdaSAM, has already been explored. We conduct several experiments on several NLP tasks, which show that AdaSAM could achieve superior performance compared with SGD, AMS, and SAMsGrad.
  arXiv  Detail & Related papers  (2023-03-01T15:12:42Z)
• Weight Update Skipping: Reducing Training Time for Artificial Neural Networks [0.30458514384586394]
  We propose a new training methodology for ANNs that exploits the observation of improvement of accuracy shows temporal variations. During such time windows, we keep updating bias which ensures the network still trains and avoids overfitting. Such a training approach virtually achieves the same accuracy with considerably less computational cost, thus lower training time.
  arXiv  Detail & Related papers  (2020-12-05T15:12:10Z)
• Human Body Model Fitting by Learned Gradient Descent [48.79414884222403]
  We propose a novel algorithm for the fitting of 3D human shape to images. We show that this algorithm is fast (avg. 120ms convergence), robust to dataset, and achieves state-of-the-art results on public evaluation datasets.
  arXiv  Detail & Related papers  (2020-08-19T14:26:47Z)
• Communication-Efficient Distributed Stochastic AUC Maximization with Deep Neural Networks [50.42141893913188]
  We study a distributed variable for large-scale AUC for a neural network as with a deep neural network. Our model requires a much less number of communication rounds and still a number of communication rounds in theory. Our experiments on several datasets show the effectiveness of our theory and also confirm our theory.
  arXiv  Detail & Related papers  (2020-05-05T18:08:23Z)
• Lipreading using Temporal Convolutional Networks [57.41253104365274]
  Current model for recognition of isolated words in-the-wild consists of a residual network and Bi-directional Gated Recurrent Unit layers. We address the limitations of this model and we propose changes which further improve its performance. Our proposed model results in an absolute improvement of 1.2% and 3.2%, respectively, in these datasets.
  arXiv  Detail & Related papers  (2020-01-23T17:49:35Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.