mlx.optimizers.Lion

目录

mlx.optimizers.Lion#

class Lion(learning_rate: float | Callable[[array], array], betas: List[float] = [0.9, 0.99], weight_decay: float = 0.0)#

Lion 优化器 [1]。

由于更新是通过符号函数计算的，因此它们的范数往往比 SGD 和 Adam 等其他优化器更大。我们建议将学习率设为 AdamW 的 3-10 倍小，权重衰减设为 AdamW 的 3-10 倍大，以保持强度 (lr * wd)。我们的 Lion 实现遵循原始论文。详细来说，

[1]: Chen, X. Symbolic Discovery of Optimization Algorithms. arXiv preprint arXiv:2302.06675.

\[\begin{split}c_{t + 1} &= \beta_1 m_t + (1 - \beta_1) g_t \\ m_{t + 1} &= \beta_2 m_t + (1 - \beta_2) g_t \\ w_{t + 1} &= w_t - \eta (\text{sign}(c_t) + \lambda w_t)\end{split}\]
参数:

  • learning_rate (float 或 callable) – 学习率 \(\eta\)。

  • betas (Tuple[float, float], 可选) – 用于计算梯度动量和更新方向的系数 \((\beta_1, \beta_2)\)。默认值: (0.9, 0.99)

  • weight_decay (float, 可选) – 权重衰减 \(\lambda\)。默认值: 0.0

方法

__init__(learning_rate[, betas, weight_decay])

apply_single(gradient, parameter, state)

执行 Lion 参数更新并将 \(m\) 存储在优化器状态中。

init_single(parameter, state)

初始化优化器状态
目录