Google Brain新提出的优化器“Lion”，作用要比Adam(W)更好

  • 论文地址：arxiv.org/abs/2302.06…
  • 代码地址：github.com/google/auto…

1 简略、内存高效、运转速度更快

与 AdamW 和各种自适应优化器需求一起保存一阶和二阶矩比较，Lion 只需求动量，将额外的内存占用减半。 这在练习大型模型和大Batch size时很有用。 例如，AdamW 需求至少 16 个 TPU V4 芯片来练习图画巨细为 224、批量巨细为 4,096 的 ViT-B/16，而 Lion 只需求8个。

另一个清楚明了的优点是，因为 Lion 的简略性，Lion 在咱们的试验中具有更快的运转时间（step/s），一般比 AdamW 和 Adafactor 提速 2-15%，详细取决于使命、代码库和硬件。

Google Brain新提出的优化器“Lion”，效果要比Adam(W)更好

2 在各种模型、使命和范畴上的优越功能

2.1 图画分类

  • Lion 在 ImageNet 上从头开始练习或在 ImageNet-21K 上预练习的各种网络模型上优于 AdamW。

Google Brain新提出的优化器“Lion”，效果要比Adam(W)更好

  • Lion 在 JFT-300M 上节省了高达 5 倍的预练习成本。

Google Brain新提出的优化器“Lion”，效果要比Adam(W)更好

  • 运用更高分辨率和 Polyak 均匀进行微调后的成果。 Lion取得的 ViT-L/16 与之前由 AdamW 练习的 ViT-H/14 成果相匹配，一起缩小了 2 倍，一起关于 ViT-G/14 在 ImageNet 进步一步达到了 90.71% 的精确率。
    Google Brain新提出的优化器“Lion”，效果要比Adam(W)更好

2.2 视觉-言语比照练习

  • 在 LiT 上，Lion 在零样本图画分类和图画文本检索方面击败了 AdamW。

Google Brain新提出的优化器“Lion”，效果要比Adam(W)更好

Google Brain新提出的优化器“Lion”，效果要比Adam(W)更好

  • 在 BASIC-L 上，Lion 实现了 88.3% 的零样本和 91.1% 的微调 ImageNet 精确率，别离超越之前的最佳成果 2% 和 0.1%。

Google Brain新提出的优化器“Lion”，效果要比Adam(W)更好

2.3 分散模型

  • 在分散模型上，Lion 在 FID 分数方面超越了 AdamW，节省了高达 2.3 倍的练习核算。 从左到右：在 ImageNet 上练习的 64×64、128×128、256×256 图画生成。

Google Brain新提出的优化器“Lion”，效果要比Adam(W)更好

2.4 言语建模

  • Lion 在执行言语建模使命时在验证困惑度（perplexity）上节省了高达 2 倍的核算量（左：在 Wiki-40B 上，右：在 PG-19 上）。 Lion 在更大的transformer上取得更大的收益。

Google Brain新提出的优化器“Lion”，效果要比Adam(W)更好

  • 与 Adafactor 比较，Lion 在练习 LLM 时取得更好的均匀上下文学习才能。

Google Brain新提出的优化器“Lion”，效果要比Adam(W)更好

  • 在 GLUE 上微调 T5 时 Lion 也更好。

Google Brain新提出的优化器“Lion”，效果要比Adam(W)更好

3 超参数和批量巨细挑选

  • Lion 很简略，与 AdamW 和 Adafactor 比较，超参数更少，因为它不需求 \epsilon 和因式分解相关的参数。 为了确保公平比较，咱们运用对数标度为 AdamW (Adafactor) 和咱们的 Lion 调整峰值学习率 lrlr 宽和耦权重衰减 \lambda。 AdamW 中 1\beta_1 和 2\beta_2 的默认值别离设置为 0.9 和 0.999，\epsilon 为 1e−81e-8，而在 Lion 中，1\beta_1 和 1\beta_1 的默认值 2\beta_2 是经进程序搜索进程发现的，别离设置为 0.9 和 0.99。 作者只调整言语使命中的那些超参数，其中 1=0.9\beta_1=0.9，2=0.99\beta_2=0.99 在 AdamW 中，1=0.95\beta_1=0.95，2=0.98\beta_2=0.98 在 Lion 中。 此外，AdamW 中的 \epsilon 设置为 1e−61e-6 而不是默认的 1e−81e-8，因为它提高了咱们试验中的稳定性，类似于 RoBERTa 中的调查成果。

  • Lion 生成的更新是元素二进制 1\pm 1，作为符号操作的成果，因而它具有比其他优化器生成的更大的范数。 依据作者的经历，Lion 的适宜学习率一般比 AdamW 小 10 倍，虽然有时小 3 倍的学习率可能体现稍好。 因为有效权重衰减为 lr∗lr * \lambda，因而用于 Lion 的 \lambda 值比 AdamW 大 10 倍，以坚持相似的强度。 例如，

    • lr=1e−4lr=1e-4, =10.0\lambda=10.0 在 Lion 和 lr=1e−3lr=1e-3, =1.0\lambda=1.0 在 ImageNet 上练习 ViT-B/16 时运用强增强。
    • Lion 中的 lr=3e−5lr=3e-5, =0.1\lambda=0.1 和 AdamW 中的 lr=3e−4lr=3e-4, =0.01\lambda=0.01 用于分散模型。
    • Lion 中的 lr=1e−4lr=1e-4、=0.01\lambda=0.01 和 Adafactor 中的 lr=1e−3lr=1e-3、=0.001\lambda=0.001 用于 7.5B 言语建模。

  • 除了峰值功能外，对超参数的敏感性和调整它们的难度关于在实践中采用优化器也很关键。 在下图中，咱们在 ImageNet 上从头开始练习 ViT-B/16 时一起更改 lrlr 和 \lambda。 热图标明，与 AdamW 比较，Lion 关于不同的超参数挑选愈加稳健。

  • 有些人可能会质疑 Lion 是否需求大批量巨细才能精确确定方向，因为标志操作会添加噪音。 为了解决这个问题，咱们运用各种批量巨细在 ImageNet 上练习 ViT-B/16 模型，一起将总练习时期坚持为 300，并结合 RandAug 和 Mixup 技术。 如下图所示，AdamW 的最佳批量巨细为 256，而 Lion 为 4,096。 这标明 Lion 的确更喜爱更大的批处理巨细，但即使运用 64 的小批处理巨细，其功能仍然坚持稳健。 此外，当批量巨细扩大到 32K 时，只需求 11K 练习过程， Lion 的精确率比 AdamW 高出 2.5%（77.9% 对 75.4%），证明了它在大批量练习环境中的有效性。

Google Brain新提出的优化器“Lion”，效果要比Adam(W)更好

左：批量巨细影响的融化试验。 Lion 比 AdamW 更喜爱更大的批次。 当咱们为 AdamW（中间）和 Lion（右）改变 lrlr 和 \lambda 时，从头开始练习的 ViT-B/16 的 ImageNet 精度。 Lion 关于不同的超参数挑选愈加稳健。

4 代码实现

"""PyTorch implementation of the Lion optimizer."""
import torch
from torch.optim.optimizer import Optimizer
class Lion(Optimizer):
  r"""Implements Lion algorithm."""
  def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0):
    """Initialize the hyperparameters.
    Args:
      params (iterable): iterable of parameters to optimize or dicts defining
        parameter groups
      lr (float, optional): learning rate (default: 1e-4)
      betas (Tuple[float, float], optional): coefficients used for computing
        running averages of gradient and its square (default: (0.9, 0.99))
      weight_decay (float, optional): weight decay coefficient (default: 0)
    """
    if not 0.0 <= lr:
      raise ValueError('Invalid learning rate: {}'.format(lr))
    if not 0.0 <= betas[0] < 1.0:
      raise ValueError('Invalid beta parameter at index 0: {}'.format(betas[0]))
    if not 0.0 <= betas[1] < 1.0:
      raise ValueError('Invalid beta parameter at index 1: {}'.format(betas[1]))
    defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
    super().__init__(params, defaults)
  @torch.no_grad()
  def step(self, closure=None):
    """Performs a single optimization step.
    Args:
      closure (callable, optional): A closure that reevaluates the model
        and returns the loss.
    Returns:
      the loss.
    """
    loss = None
    if closure is not None:
      with torch.enable_grad():
        loss = closure()
    for group in self.param_groups:
      for p in group['params']:
        if p.grad is None:
          continue
        # Perform stepweight decay
        p.data.mul_(1 - group['lr'] * group['weight_decay'])
        grad = p.grad
        state = self.state[p]
        # State initialization
        if len(state) == 0:
          # Exponential moving average of gradient values
          state['exp_avg'] = torch.zeros_like(p)
        exp_avg = state['exp_avg']
        beta1, beta2 = group['betas']
        # Weight update
        update = exp_avg * beta1 + grad * (1 - beta1)
        p.add_(torch.sign(update), alpha=-group['lr'])
        # Decay the momentum running average coefficient
        exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2)
    return loss

5 参考资料

  • github.com/google/auto…