mlx.optimizers.Adafactor

目录

mlx.optimizers.Adafactor#

class Adafactor(learning_rate: float | Callable[[array], array] | None = None, eps: Tuple[float, float] = (1e-30, 0.001), clip_threshold: float = 1.0, decay_rate: float = -0.8, beta_1: float | None = None, weight_decay: float = 0.0, scale_parameter: bool = True, relative_step: bool = True, warmup_init: bool = False)#

Adafactor 优化器。

我们的 Adafactor 实现遵循原始论文：Adafactor：具有亚线性内存成本的自适应学习率

参数:

  • learning_rate (float 或 callable, 可选) – 学习率。默认值：None。

  • eps (tuple(float, float), 可选) – 添加到梯度平方中的第一项 \(\epsilon_1\) 用于提高数值稳定性；第二项 \(\epsilon_2\) 在 parameter_scale 设置为 True 时用于参数缩放。默认值：(1e-30, 0.001)。

  • clip_threshold (float, 可选) – 将未缩放的更新限制在 clip_threshold。默认值：1.0。

  • decay_rate (float, 可选) – 平方梯度滑动平均的系数。默认值：-0.8。

  • beta_1 (float, 可选) – 如果设置为大于零的值，则将使用一阶矩。默认值：None。

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

  • scale_parameter (bool, 可选) – 如果设置为 True，学习率将按 \(\max(\epsilon_1, \text{RMS}(w_{t-1}))\) 进行缩放。默认值：True。

  • relative_step (bool, 可选) – 如果设置为 True，则 learning_rate 将被忽略，并计算相对步长。默认值：True。

  • warmup_init (bool, 可选) – 如果设置为 True，则相对步长将根据当前步数计算。默认值：False。

方法

__init__([learning_rate, eps, ...])

apply_single(gradient, parameter, state)

执行 Adafactor 参数和状态更新。

init_single(parameter, state)

初始化优化器状态
目录