Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
AF-Muon позволил обучать модели с общей таблицей токенов без вспомогательного AdamW. Хотя препринт не рецензирован и числа получили сами авторы, AF-Muon сократил память под состояние оптимизатора примерно на 20% относительно Hybrid Muon, прибавив около 1% ко времени шага. Схема освобождает память под модель или пакет, не требуя менять матричную часть Muon.
В обычной схеме Muon обновляет скрытые матрицы, а эмбеддинги, выходной слой, параметры нормализации и смещения передаёт AdamW. Сложнее всего устроена общая таблица токенов (tied embeddings): она одновременно выдаёт векторы входных токенов и превращает скрытое состояние в вероятности выходных токенов. Поэтому в ней смешиваются разреженные градиенты от входа и плотные градиенты от классификатора.
AF-Muon сохраняет обновление Muon для скрытых матриц, но разделяет остальные параметры по роли. Одномерные параметры получают обновление, нормированное по среднеквадратичному значению, а общая таблица — отдельное обновление с ограничением на каждую строку и координату. Ограничение сохраняет различия между крупными и мелкими компонентами градиента, но не позволяет нескольким координатам забрать почти всю величину шага.
Все классы параметров хранят только накопленный первый момент градиента. Второй момент, ради которого вспомогательная ветка AdamW держит дополнительный массив состояния, больше не нужен. При этом простое удаление AdamW не объясняет результат: вариант с конечным ограничением для общей таблицы оказался лучше крайнего варианта SCION-style Sign, который выравнивает величину всех ненулевых координат.
Метод проверили в девяти конфигурациях: на декодерных языковых моделях от 124 млн до 1 млрд параметров, общем кодировщике-декодере, разреженной смеси экспертов, модели изображений и модели белковых последовательностей. При одинаковых бюджетах AF-Muon показал меньшие средние потери на проверочной выборке и меньшую перплексию, чем Hybrid Muon и SCION-style Sign, во всех конфигурациях. Вывод относится к моделям с общей входной и выходной таблицей токенов: для несвязанных таблиц работа такого результата не устанавливает.
Источники
Иллюстрация: рисунок из статьи «AF-Muon: An AdamW-Free Muon Optimizer for Tied-Embedding Models», Arash Lagzian, Paniz Halvachi, Junming Zhang и др., CC BY 4.0
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
