🔥 Lion: EvoLved Sign Momentum - новый оптимайзер от Гугла

Google нашел новый алгоритм оптимизации при помощи Symbolic Search (завтра объясню, что это такое). Новый оптимайзер тупо мониторит momentum, а размер финального апдейта весов считается как sign(grad (1-β1) + β1 momentum). То есть апдейт для каждого элемента весов равен либо -1 * LR либо +1 * LR!

train(weight, gradient, momentum, lr):
update = interp(gradient, momentum, β1)
update = sign(update)
momentum = interp(gradient, momentum, β2)
// Тут я пропустил weight decay для простоты
update = update * lr
return update, momentum

Просто волосы дыбом от тупости Lion, но он дает серьезные улучшения в скорости сходимости, да и в найденных локальных минимумах. Рвет и обычную классификацию и языковые модели и диффузионные.

Чем больше батч, тем лучше работает Lion. Еще он более устойчив к различным выборам гиперпараметров по сравнению с AdamW.

Уже попробовал на свои задачах, вроде работает!

❱❱ Код на PyTorch, JAX и TF

@ai_newz