adv-lion

Advanced Lion Optimizer

This repository provides an enhanced implementation of the Lion optimizer, incorporating several state-of-the-art techniques to improve performance, stability, and memory efficiency. It includes two base variants: the original Lion and D-Adapt-Lion.

---

Features

1. Fused Backward Pass

Reduces memory overhead by hooking gradients as they become available during the backward pass, eliminating the need to store them explicitly via the step_parameter method.

---

2. Stochastic Rounding for BF16 Training

Achieves FP32-level performance in BF16 training by implementing stochastic rounding for the final update. This allows for faster training with lower-precision data types without sacrificing accuracy.

• References:
  • "Revisiting BFloat16 Training"
  • "Stochastic Rounding for LLM Training: Theory and Practice"

---

3. Gradient Orthogonalization

Improves model generalization and enhances numerical stability by modifying gradients to be orthogonal to the parameters.

• Reference: "Grokking at the Edge of Numerical Stability"

---

4. Variance Reduction

Theoretically accelerates the convergence speed of Lion by 33.33% while making training more stable in noisy, small-batch environments. The main trade-off is the requirement of an additional state to store gradients from the previous step.

• Reference: "Convergence Analysis of the Lion Optimizer in Centralized and Distributed Settings"

---

5. Cautious Lion Variant

Includes the "Cautious" variant of Lion, an approach introduced to refine the optimization process and improve training outcomes.

• Reference: "Cautious Optimizers: Improving Training with One Line of Code"

---

6. Per-Parameter Gradient Norm Clipping

Enhances training stability by applying gradient norm clipping on a per-parameter basis, preventing erratic updates from large gradients.

• Reference: "Lions and Muons: Optimization via Stochastic Frank-Wolfe" (The paper uses a clipping value of 4-5).