lion-pytorch

<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/565dde9c-2502-4df4-9691-de7314dde220.png" width="500px"></img>

🦁 Lion - Pytorch

<a href="https://arxiv.org/abs/2302.06675">🦁 Lion</a>, EvoLved Sign Momentum，是谷歌大脑发现的新优化器，据称比Adam(w)更好，现已在Pytorch上实现。这几乎是从<a href="https://github.com/google/automl/blob/master/lion/lion_pytorch.py">这里</a>复制过来的，仅做了一些小改动。

它非常简单，如果真的有效，我们不妨尽快让大家使用它来训练一些优秀的模型。

使用说明

• 学习率和权重衰减：作者在第5节中写道——根据我们的经验，Lion的合适学习率通常比AdamW小3-10倍。由于有效权重衰减为学习率 * λ，为了保持相似的强度，Lion使用的分离权重衰减λ值比AdamW大3-10倍。 学习率计划中的初始值、峰值和结束值应当同时按与AdamW比较的相同比例进行变化，由研究人员证实。

• 学习率计划：作者在论文中对Lion和AdamW使用了相同的学习率计划。不过，他们注意到，使用余弦衰减计划训练ViT时，相较于倒平方根计划，获得了更大的增益。

• β1和β2：作者在第5节中写道——AdamW中的默认β1和β2值分别设为0.9和0.999，ε为1e−8，而在Lion中，默认的β1和β2值通过程序搜索过程发现，分别设为0.9和0.99。 类似于人们通过将β2减少到0.99或更小并将ε增加到1e-6来提高AdamW的稳定性，作者建议在Lion中使用 β1=0.95, β2=0.98 也有助于缓解训练中的不稳定性。通过研究人员证实。

更新

• 更新：在我的本地enwik8自回归语言建模中似乎有效。

• 更新2：<a href="https://api.wandb.ai/links/lucidrains/d4v6c8sl">实验</a>，如果学习率保持不变，效果似乎比Adam差得多。

• 更新3：将学习率除以3，看到了比Adam更好的早期结果。也许Adam在近十年后终于被取代了。

• 更新4：使用论文中的10倍较小的学习率经验法则导致最差的运行。所以我猜还需要一些调试。

之前更新的总结：如<a href="https://api.wandb.ai/links/lucidrains/d4v6c8sl">实验</a>所示，Lion使用3倍较小的学习率击败了Adam。仍然需要一些调试，因为10倍较小的学习率会导致较差的结果。

• 更新5：到目前为止，听到的关于语言建模的结果都是积极的，只要做得正确。此外，还听到了对大型文本到图像训练的积极结果，虽然需要一些调试。负面结果似乎与论文评估之外的问题和架构有关——RL、前馈网络、与LSTM+卷积等组合的奇怪混合架构。负面的零散数据也证实了这种技术对批次大小、数据量/增强量敏感。待确定最佳的学习率计划是什么，以及降温是否影响结果。此外，有一个有趣的正面结果是在open-clip上，但当模型规模扩大时结果变为负面（可能是可解决的）。

• 更新6：open clip问题已被作者通过设定更高的初始温度解决。

• 更新7：仅推荐在大批次（64或以上）的设置中使用此优化器。

安装

$ pip install lion-pytorch

使用

# 玩具模型

import torch
from torch import nn

model = nn.Linear(10, 1)

# 导入Lion并用参数实例化

from lion_pytorch import Lion

opt = Lion(model.parameters(), lr=1e-4, weight_decay=1e-2)

# 前向和后向

loss = model(torch.randn(10))
loss.backward()

# 优化器步骤

opt.step()
opt.zero_grad()

要使用融合的内核更新参数，首先 pip install triton -U --pre，然后

opt = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    use_triton=True # 将其设为True以使用带Triton语言（Tillet et al）的CUDA内核
)

感谢

• 感谢 <a href="https://stability.ai/">Stability.ai</a> 慷慨赞助以开展和开源前沿人工智能研究

引用

@misc{https://doi.org/10.48550/arxiv.2302.06675,
    url     = {https://arxiv.org/abs/2302.06675},
    author  = {Chen, Xiangning and Liang, Chen and Huang, Da and Real, Esteban and Wang, Kaiyuan and Liu, Yao and Pham, Hieu and Dong, Xuanyi and Luong, Thang and Hsieh, Cho-Jui and Lu, Yifeng and Le, Quoc V.},
    title   = {Symbolic Discovery of Optimization Algorithms},
    publisher = {arXiv},
    year = {2023}
}

@article{Tillet2019TritonAI,
    title   = {Triton: an intermediate language and compiler for tiled neural network computations},
    author  = {Philippe Tillet and H. Kung and D. Cox},
    journal = {Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages},
    year    = {2019}
}

@misc{Schaipp2024,
    author  = {Fabian Schaipp},
    url     = {https://fabian-sp.github.io/posts/2024/02/decoupling/}
}