SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training

• URL: http://arxiv.org/abs/2412.13148v2
• Date: Mon, 23 Dec 2024 10:46:21 GMT
• Title: SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
• Authors: Chao Ma, Wenbo Gong, Meyer Scetbon, Edward Meeds,
• Abstract summary: Gradient Descent (SGD) is a stateless, scalability as it does not track state variables during training. In this work, we show that pre-processing SGD in a stateless manner can achieve the same performance as the Adam for training LLM. We show that normalization stabilizes gradient, and counteracts the local curvature of the loss landscape. This results in SWAN (SGD with Whitening And Normalization), aapprox that eliminates the need to store any states.
• Score: 16.037614012166063
• License:
• Abstract: Adaptive optimizers such as Adam (Kingma & Ba, 2015) have been central to the success of large language models. However, they often require to maintain optimizer states throughout training, which can result in memory requirements several times greater than the model footprint. This overhead imposes constraints on scalability and computational efficiency. Stochastic Gradient Descent (SGD), in contrast, is a stateless optimizer, as it does not track state variables during training. Consequently, it achieves optimal memory efficiency. However, its capability in LLM training is limited (Zhao et al., 2024b). In this work, we show that pre-processing SGD in a stateless manner can achieve the same performance as the Adam optimizer for LLM training, while drastically reducing the memory cost. Specifically, we propose to pre-process the instantaneous stochastic gradients using normalization and whitening. We show that normalization stabilizes gradient distributions, and whitening counteracts the local curvature of the loss landscape. This results in SWAN (SGD with Whitening And Normalization), a stochastic optimizer that eliminates the need to store any optimizer states. Empirically, SWAN has the same memory footprint as SGD, achieving $\approx 50\%$ reduction on total end-to-end memory compared to Adam. In language modeling tasks, SWAN demonstrates comparable or even better performance than Adam: when pre-training the LLaMA model with 350M and 1.3B parameters, SWAN achieves a 2x speedup by reaching the same evaluation perplexity using half as many tokens.

Related papers

• Gradient Multi-Normalization for Stateless and Scalable LLM Training [16.037614012166063]
  Training large language models (LLMs) typically relies on adaptives like Adam which store additional state information to accelerate convergence but incur significant memory overhead. Recent efforts, such as SWAN (Ma et al., 2024) address this by eliminating the need for states while achieving performance comparable to Adam via a multi-step preprocessing procedure applied to instantaneous gradients. We introduce a novel framework for designing stateless gradients that normalizes gradients according to multiple norms. Experiments on pre-training LLaMA models with up to 1 billion parameters demonstrate a 3X speedup over Adam with significantly reduced memory requirements, outperforming other memory-efficient baseline
  arXiv  Detail & Related papers  (2025-02-10T18:09:53Z)
• APOLLO: SGD-like Memory, AdamW-level Performance [61.53444035835778]
  Large language models (LLMs) are notoriously memory-intensive during training. Various memory-efficient Scals have been proposed to reduce memory usage. They face critical challenges: (i) costly SVD operations; (ii) significant performance trade-offs compared to AdamW; and (iii) still substantial memory overhead to maintain competitive performance.
  arXiv  Detail & Related papers  (2024-12-06T18:55:34Z)
• COAP: Memory-Efficient Training with Correlation-Aware Gradient Projection [11.655821671462427]
  We present COAP, a memory-efficient method that minimizes computational overhead while maintaining training performance. For LLaMA-1B, it reduces memory by 61% with only 2% additional time cost, achieving the same PPL as AdamW. With 8-bit quantization, COAP cuts memory by 81% and 4x speedup over GaLore for LLaVA-v1.5-7B fine-tuning, while delivering higher accuracy.
  arXiv  Detail & Related papers  (2024-11-26T03:50:52Z)
• FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training [51.39495282347475]
  We introduce $textttFRUGAL$ ($textbfF$ull-$textbfR$ank $textbfU$pdates with $textbfG$r$textbfA$dient sp$textbfL$itting, a new memory-efficient optimization framework. Our framework can be integrated with various low-rank update selection techniques, including GaLore and BAdam.
  arXiv  Detail & Related papers  (2024-11-12T14:41:07Z)
• LDAdam: Adaptive Optimization from Low-Dimensional Gradient Statistics [37.21593513802284]
  We introduce LDAdam, a memory-efficient gradient for training large models. We show that LDAdam allows for accurate and efficient fine-tuning and pre-training of language models.
  arXiv  Detail & Related papers  (2024-10-21T15:31:06Z)
• Simultaneous Computation and Memory Efficient Zeroth-Order Optimizer for Fine-Tuning Large Language Models [33.911521719528686]
  Fine-tuning is powerful for adapting large language models to downstream tasks, but it often results in huge memory usages. A promising approach is using Zeroth-Order (ZO) gradients, which estimates to replace First-Order (FO) gradients. We introduce a novel layer-wise sparse computation and memory efficient ZO, named LeZO.
  arXiv  Detail & Related papers  (2024-10-13T12:47:37Z)
• Zeroth-Order Fine-Tuning of LLMs in Random Subspaces [66.27334633749734]
  As language models grow in size, memory demands for backpropagation increase. Zeroth-order (ZOZO) optimization methods offer a memory-efficient alternative. We show that SubZero enhances fine-tuning and achieves faster results compared to standard ZOZO approaches.
  arXiv  Detail & Related papers  (2024-10-11T17:01:43Z)
• Bypass Back-propagation: Optimization-based Structural Pruning for Large Language Models via Policy Gradient [57.9629676017527]
  We propose an optimization-based structural pruning on Large-Language Models. We learn the pruning masks in a probabilistic space directly by optimizing the loss of the pruned model. Our method operates for 2.7 hours with around 35GB memory for the 13B models on a single A100 GPU.
  arXiv  Detail & Related papers  (2024-06-15T09:31:03Z)
• AdaLomo: Low-memory Optimization with Adaptive Learning Rate [59.64965955386855]
  We introduce low-memory optimization with adaptive learning rate (AdaLomo) for large language models. AdaLomo results on par with AdamW, while significantly reducing memory requirements, thereby lowering the hardware barrier to training large language models.
  arXiv  Detail & Related papers  (2023-10-16T09:04:28Z)
• Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model [89.8764435351222]
  We propose a new family of unbiased estimators called WTA-CRS, for matrix production with reduced variance. Our work provides both theoretical and experimental evidence that, in the context of tuning transformers, our proposed estimators exhibit lower variance compared to existing ones.
  arXiv  Detail & Related papers  (2023-05-24T15:52:08Z)

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.