Muon, Keller Jordan 2024

• Record holder for nanoGPT (source) TODO

• Relation to shampoo: (source)

  A connection between Muon and Shampoo (https://arxiv.org/abs/1802.09568) is that in the momentum=0 case, running Shampoo with the preconditioner updated every step and with no accumulation also yields the nearest orthogonal matrix as its update, although this has a slower runtime.

  The nearest orthogonal matrix is also equivalent to:

  1. The semi-unitary factor of the update's polar decomposition.
  2. UV^T where USV^T is the update's SVD.

  Using a rescaled version of (2) was proposed in prior work (

  https://ieeexplore.ieee.org/document/7347351…).

SOAP, Harvard 2024

• SOAP: Improving and Stabilizing Shampoo using Adam

• Algorithm

  • Infrequent Operations (every f steps)

    • Update preconditioner matrices:
      • $L_t = \beta_2 L_{t-1} + (1-\beta_2)G_tG_t^T$
      • $R_t = \beta_2 R_{t-1} + (1-\beta_2)G_t^TG_t$
    • Compute rotation matrices (eigenvectors):
      • $L_t = Q_L\Lambda_LQ_L^T$
      • $R_t = Q_R\Lambda_RQ_R^T$

    # Compute preconditioner matrices (like Shampoo)
    L = β2 * L + (1-β2) * GG^T
    R = β2 * R + (1-β2) * G^TG
    
    # Find rotation matrices (eigenvectors)
    QL = find_eigenvectors(L)
    QR = find_eigenvectors(R)
    

  • Every Step Operations:

    # Rotate gradient to new coordinate system
    G_rotated = QL^T @ G @ QR
    
    # Run regular Adam steps in rotated space
    m_rotated = β1 * m_rotated + (1-β1) * G_rotated  # momentum
    v_rotated = β2 * v_rotated + (1-β2) * G_rotated²  # second moment
    update = m_rotated / sqrt(v_rotated)
    
    # Rotate update back to original space
    final_update = QL @ update @ QR^T
    
    # Final weight update
    weight -= lr * final_update
    

• Vs Shampoo

  • Shampoo Also Rotates, BUT:
    • Shampoo does the rotation and preconditioning in one step
    • When Shampoo updates its preconditioner (L and R matrices), it has to recompute everything
  • Why SOAP is More Robust:
    • If Shampoo's update frequency is too low, performance degrades badly
    • SOAP degrades more gracefully because Adam keeps adapting between rotation updates
    • This is why SOAP can use less frequent updates (cheaper) while maintaining good performance
  • Also SOAP needs fewer hyperparameters - much of the complexity is handled by Adam's adaptation in the rotated space, rather than needing to carefully tune the preconditioning

Overview:

1. Background and Problem:
   • Shampoo is a higher-order optimization method that has shown better performance than Adam in deep learning
   • However, Shampoo has drawbacks: additional hyperparameters and computational overhead compared to Adam
2. Main Contribution - SOAP Algorithm:
   • The authors establish that Shampoo (with 1/2 power) is equivalent to running Adafactor in the eigenbasis of Shampoo's preconditioner
   • Based on this insight, they develop SOAP, which runs Adam in Shampoo's eigenbasis
   • SOAP retains the benefits of both methods while addressing Shampoo's limitations
3. Key Results:
   • For large language model training (360M and 660M parameter models):
     • 40% reduction in number of iterations compared to AdamW
     • 35% reduction in wall clock time compared to AdamW
     • 20% improvement in both metrics compared to Shampoo
   • SOAP shows better robustness to preconditioning frequency compared to Shampoo
   • SOAP requires only one additional hyperparameter compared to AdamW (preconditioning frequency)
4. Implementation Details:
   • The algorithm computes eigenvectors less frequently to improve efficiency
   • Maintains running averages of second moments in the current coordinate basis
   • Uses power iteration and QR decomposition for eigenvector computation
   • Special handling for large layers (like first/last layers in transformers)
5. Future Work:
   • Exploring lower precision computations for better efficiency
   • Developing improved distributed implementations
   • Testing performance on other domains beyond language models

The significance of this work is that it combines the benefits of first-order methods (Adam) and second-order methods (Shampoo) in a principled way, resulting in an optimizer that is both more efficient and easier to tune than existing approaches. The empirical results on large language model training demonstrate substantial practical benefits.

Shampoo

• Shampoo optimizer