Lion优化器: Google Brain发现的新型神经网络优化算法

Lion优化器:突破性的神经网络训练工具

在深度学习领域,优化算法的选择对模型训练效果至关重要。近日,Google Brain团队通过遗传算法发现了一种名为Lion(EvoLved SIgn MOmeNtum)的新型优化器,引起了人工智能研究界的广泛关注。本文将详细介绍Lion优化器的特点、使用方法及其在实际应用中的表现。

Lion优化器简介

Lion优化器是Google Brain团队在符号程序搜索过程中发现的一种新型优化算法。它的全称是EvoLved SIgn MOmeNtum,意为"进化的符号动量"。根据研究团队的说法,Lion在多项任务中的表现超过了目前广泛使用的Adam(w)优化器。

Lion优化器与其他优化器的性能对比

Lion优化器的主要特点

学习率设置:Lion的学习率通常需要比AdamW小3-10倍。这意味着如果你之前使用AdamW时的学习率是1e-3,那么使用Lion时可能需要将学习率调整到1e-4左右。
权重衰减:由于有效权重衰减是学习率与λ的乘积,因此Lion使用的解耦权重衰减λ值通常需要比AdamW大3-10倍,以保持相似的强度。
学习率调度:研究人员发现,使用余弦衰减调度训练ViT模型时,Lion的性能增益比使用倒数平方根调度更大。
β1和β2参数:Lion的默认β1和β2值分别为0.9和0.99,这是通过程序搜索过程发现的。为了提高训练稳定性,可以尝试使用β1=0.95,β2=0.98的设置。

Lion优化器的使用方法

要在PyTorch中使用Lion优化器,首先需要安装lion-pytorch库:

pip install lion-pytorch

然后,可以按照以下方式在代码中使用Lion优化器:

import torch
from torch import nn
from lion_pytorch import Lion

# 创建模型
model = nn.Linear(10, 1)

# 实例化Lion优化器
opt = Lion(model.parameters(), lr=1e-4, weight_decay=1e-2)

# 前向传播和反向传播
loss = model(torch.randn(10))
loss.backward()

# 优化器步骤
opt.step()
opt.zero_grad()

如果想使用融合内核来更新参数,可以先安装triton库,然后在实例化Lion优化器时设置use_triton=True:

pip install triton -U --pre

opt = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    use_triton=True  # 使用CUDA内核与Triton语言
)

Lion优化器的实际表现

根据多位研究者的实验和反馈,Lion优化器在实际应用中表现出以下特点:

在语言建模任务中,当参数设置合适时,Lion通常能够取得优于Adam的效果。
在大规模文本到图像训练中也有不错的表现,但可能需要一些调优工作。
Lion对批量大小和数据量/增强的敏感度较高,建议在批量大小大于或等于64的情况下使用。
在强化学习、前馈网络等论文中未评估的问题和架构上,Lion的表现可能不如预期。
在某些情况下,Lion能够以更少的训练步骤达到与Adam相同的损失水平。例如,在一项小规模评估中,Lion仅用了Adam 62.5%的步数就达到了相似的训练损失。