🔥 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