title: No More Adam: Learning Rate Scaling at Initialization is All You Need
summary: In this work, we question the necessity of adaptive gradient methods for
  training deep neural networks. SGD-SaI is a simple yet effective enhancement to
  stochastic gradient descent with momentum (SGDM). SGD-SaI performs learning
  rate Scaling at Initialization (SaI) to distinct parameter groups, guided by
  their respective gradient signal-to-noise ratios (g-SNR). By adjusting learning
  rates without relying on adaptive second-order momentum, SGD-SaI helps prevent
  training imbalances from the very first iteration and cuts the optimizer's
  memory usage by half compared to AdamW. Despite its simplicity and efficiency,
  SGD-SaI consistently matches or outperforms AdamW in training a variety of
  Transformer-based tasks, effectively overcoming a long-standing challenge of
  using SGD for training Transformers. SGD-SaI excels in ImageNet-1K
  classification with Vision Transformers(ViT) and GPT-2 pretraining for large
  language models (LLMs, transformer decoder-only), demonstrating robustness to
  hyperparameter variations and practicality for diverse applications. We further
  tested its robustness on tasks like LoRA fine-tuning for LLMs and diffusion
  models, where it consistently outperforms state-of-the-art optimizers. From a
  memory efficiency perspective, SGD-SaI achieves substantial memory savings for
  optimizer states, reducing memory usage by 5.93 GB for GPT-2 (1.5B parameters)
  and 25.15 GB for Llama2-7B compared to AdamW in full-precision training
  settings.
paper_authors._id: 6761c5e3338963d15b2649e3
paper_authors.name: Minghao Xu
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-12-20T09:34:54.856Z
paper_authors._id: 667d25957a14c06257258fdb
paper_authors.avatarUrl: /avatars/49066443382ad3b7c08bb2fb18b6e136.svg
paper_authors.fullname: Minghao Xu
paper_authors.type: user
paper_authors.user: ChrisAllenMing
paper_authors._id: 6761c5e3338963d15b2649e4
paper_authors.name: Lichuan Xiang
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-12-19T18:03:44.969Z
paper_authors._id: 6639c5dd0419b50216e745ed
paper_authors.avatarUrl: /avatars/7d2342f996240e0a5c2a74aa3fcda6fa.svg
paper_authors.fullname: Lichuan
paper_authors.type: user
paper_authors.user: CodexXiang
paper_authors._id: 6761c5e3338963d15b2649e5
paper_authors.name: Xu Cai
paper_authors._id: 6761c5e3338963d15b2649e6
paper_authors.name: Hongkai Wen
paper_title: No More Adam: Learning Rate Scaling at Initialization is All You Need
paper_summary: In this work, we question the necessity of adaptive gradient methods for
  training deep neural networks. SGD-SaI is a simple yet effective enhancement to
  stochastic gradient descent with momentum (SGDM). SGD-SaI performs learning
  rate Scaling at Initialization (SaI) to distinct parameter groups, guided by
  their respective gradient signal-to-noise ratios (g-SNR). By adjusting learning
  rates without relying on adaptive second-order momentum, SGD-SaI helps prevent
  training imbalances from the very first iteration and cuts the optimizer's
  memory usage by half compared to AdamW. Despite its simplicity and efficiency,
  SGD-SaI consistently matches or outperforms AdamW in training a variety of
  Transformer-based tasks, effectively overcoming a long-standing challenge of
  using SGD for training Transformers. SGD-SaI excels in ImageNet-1K
  classification with Vision Transformers(ViT) and GPT-2 pretraining for large
  language models (LLMs, transformer decoder-only), demonstrating robustness to
  hyperparameter variations and practicality for diverse applications. We further
  tested its robustness on tasks like LoRA fine-tuning for LLMs and diffusion
  models, where it consistently outperforms state-of-the-art optimizers. From a
  memory efficiency perspective, SGD-SaI achieves substantial memory savings for
  optimizer states, reducing memory usage by 5.93 GB for GPT-2 (1.5B parameters)
  and 25.15 GB for Llama2-7B compared to AdamW in full-precision training
  settings.
paper_ai_summary: SGD-SaI, an enhancement to stochastic gradient descent with momentum, outperforms AdamW in training Transformer-based tasks by scaling learning rates based on gradient signal-to-noise ratios and reducing memory usage.
paper_ai_keywords: SGDM SGD-SaI learning rate Scaling at Initialization SaI gradient signal-to-noise ratios g-SNR adaptive gradient methods AdamW Transformer-based tasks Vision Transformers ViT GPT-2 large language models LLMs LoRA fine-tuning diffusion models memory efficiency optimizer states