Tricks to reduce instabilities

Independent weight decay

• Parameterizing weight decay independently of learning rate reduces LR sensitivity
  • recommended by authors but not applied in practice in Pytorch

Formulation

• For parameters $\theta$, let $\Delta =v/(\sqrt{u}+ϵ)$ denote the AdamW update without learning rate or weight decay.
  • For weight decay coefficient $\lambda$, max learning rate $\eta$, and schedule $s_t\in [0,1]$,
  • Original authors recommend the update $\theta \leftarrow \theta -s_t(\eta \Delta -\lambda \theta )$,
    • we refer to as independent decay.
  • Default implementation in PyTorch or Optax applies the update $\theta \leftarrow \theta -s_t\eta (\Delta -\lambda \theta )$
    • $\eta$ now scales the weigth decay term.

Attention logit growth / QK-layernorm

• Summary: $Attention(X)=\mathrm{softmax}[\frac{1}{\sqrt{d}}\mathrm{LN}(X{W}^Q)\mathrm{LN}(X{W}^K{)}^T]$
• Why ?
  • instability was observed for large models. It was caused by extremely large values in attention logits, which lead to (almost one-hot) attention weights with near-zero entropy.
  • attention logits: $z_{ij}=⟨q_i,k_j⟩/\sqrt{d_h}$
  • models with qk-layernorm exhibit considerably lower LR sensitivity and train to low loss at high learning rates

z-loss

• Another instability, when training large models, is divergence in the output logits from the log probabilities [6]. Let $y$ denote the model’s output logits, which are used to compute class probabilities $p_i$ via a softmax $p_i=e^{y_i}/Z$ where $Z={\sum }_j{e}^{y_j}$.
• This instability occurs when the logits diverge and become very negative.
• In contrast to the attention logit growth instability, this divergence occurs towards the end of training.
• The mitigation proposed by Chowdh- ery et al. [6] is to encourage log Z to remain close to zero i.e. $Z\approx 1$.
• To do so, they add an auxiliary loss ${\log }^{2}Z$, referred to as z-loss, with a coefficient of $1e-4$.