Исследователи представили MALT — оптимизатор, расширяющий возможности метода Muon для предобучения языковых моделей. В отличие от оригинального алгоритма, MALT учитывает кривизну ландшафта функции потерь через диагональное предобусловливание. Это позволяет эффективнее справляться с анизотропией градиентов и ускорять сходимость при обучении нейросетей, сохраняя при этом легковесность вычислений, характерную для Muon.
Метод Muon, ставший популярной альтернативой AdamW, использует итерации Ньютона-Шульца для ортогонализации матриц момента. Однако его зависимость от геометрии кривизны оставалась «слепой зоной». MALT устраняет этот пробел, интегрируя механизмы учета кривизны, что делает процесс обучения более стабильным и менее чувствительным к резким изменениям ландшафта потерь.
Разработка направлена на оптимизацию вычислительных затрат при сохранении высокой точности моделей. Использование диагонального предобусловливания позволяет алгоритму лучше адаптироваться к особенностям весов в глубоких архитектурах, что критически важно для масштабного предобучения современных трансформеров.
Ключевые факты
- MALT расшифровывается как Muon Augmented with Lightweight curvature-aware preconditioning.
- Алгоритм дополняет метод Muon, который использует итерации Ньютона-Шульца для ортогонализации матриц.
- Основное нововведение заключается в интеграции диагонального предобусловливания для учета кривизны ландшафта потерь.
- Метод направлен на снижение чувствительности к анизотропии кривизны, которая может замедлять обучение в стандартных реализациях Muon.