Extending $μ$P: Spectral Conditions for Feature Learning Across Optimizers

TL;DR

A novel framework to derive $\mu$P for a broader class of optimizers, including AdamW, ADOPT, LAMB, Sophia, Shampoo and Muon, is proposed and empirical insights into depth-scaling parameterization for these optimizers are provided.

Abstract

Several variations of adaptive first-order and second-order optimization methods have been proposed to accelerate and scale the training of large language models. The performance of these optimization routines is highly sensitive to the choice of hyperparameters (HPs), which are computationally expensive to tune for large-scale models. Maximal update parameterization $(μ$P$)$ is a set of scaling rules which aims to make the optimal HPs independent of the model size, thereby allowing the HPs tuned on a smaller (computationally cheaper) model to be transferred to train a larger, target model. Despite promising results for SGD and Adam, deriving $μ$P for other optimizers is challenging because the underlying tensor programming approach is difficult to grasp. Building on recent work that introduced spectral conditions as an alternative to tensor programs, we propose a novel framework to derive $μ$P for a broader class of optimizers, including AdamW, ADOPT, LAMB, Sophia, Shampoo and Muon. We implement our $μ$P derivations on multiple benchmark models and demonstrate zero-shot learning rate transfer across increasing model width for the above optimizers. Further, we provide empirical insights into depth-scaling parameterization for these optimizers.

Figures (17)

• Figure 1: $\mu$P for Sophia (trained on Llama2) - Coordinate check plots for the word embedding and output logits layers (left); Zero-shot learning rate transfer across increasing model width (middle); Decreasing training loss with increasing model width (right).
• Figure 2: (NanoGPT) Mean validation loss for increasing model width and different learning rates across four optimizers: ADOPT (top left), LAMB (top right), Sophia (bottom left), and Shampoo (bottom right). The plots demonstrate zero-shot learning rate transfer under $\mu$P (Table \ref{['tab:MuP_table']}).
• Figure 3: (Llama2) Validation loss for increasing model width and different learning rates across three optimizers: AdamW (left), ADOPT (middle), and LAMB (right). The plots demonstrate zero-shot learning rate transfer under $\mu$P (Table \ref{['tab:MuP_table']}).
• Figure 4: (Llama2 model) AdamW optimizer - Coordinate check plots under standard parameterization (top left) and under $\mu$P (bottom left) for the word embedding and output logits layers; Decreasing training loss with increasing model width under $\mu$P (right).
• Figure 5: $\mu$P for Muon (trained on Llama2) - Coordinate check plots for the word embedding and output logits layers (left); Zero-shot learning rate transfer across increasing model width (middle); Decreasing training loss with increasing model width (right).

Theorems & Definitions (3)

• Remark 1
• Remark 2
• Remark 3