Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Память, которую Adam занимает при обучении LLM, удалось сократить без заметного изменения качества и хода обучения. Работа Beihang University вышла как нерецензированный препринт, поэтому все числа получены самими авторами: в опыте с LLaMA учтённая память параметров и оптимизатора уменьшилась на 29%. Для команд это вариант сохранить знакомый Adam, а не переходить на оптимизатор с другой логикой и заново подбирать режим обучения.
AdamPM хранит одну оценку масштаба для соседних параметров
Adam поддерживает для каждого параметра два дополнительных состояния. Первое — сглаженный градиент, который в основном задаёт направление следующего шага. Второе — сглаженный квадрат градиента, по которому оптимизатор регулирует размер шага отдельно для каждого параметра.
Из-за этих состояний Adam хранит ещё два массива размером с модель. При обучении крупной LLM они конкурируют за память с параметрами, градиентами и промежуточными результатами вычислений.
AdamPM оставляет первое состояние без изменений, а соседние элементы второго объединяет в группы. Оптимизатор усредняет квадраты градиентов внутри каждой группы, обновляет одно общее состояние и временно распространяет его обратно на все элементы перед изменением весов.
Порядок операций здесь определяет экономию. Усреднение происходит до обновления экспоненциального скользящего среднего, поэтому полное второе состояние не нужно хранить между шагами. При размере группы s оно занимает 1/s памяти исходного массива.
По умолчанию группы идут вдоль строк матрицы весов. Если размер не делится на длину группы, AdamPM дополняет последнюю группу её средним значением. Тот же механизм применяется к матрицам внимания, полносвязным слоям, таблице токенов и одномерным параметрам нормализации.
Скорость обучения, коэффициенты затухания и темп обучения остаются такими же, как у Adam. Это отличает AdamPM от методов с низкоранговыми проекциями или изменённым правилом шага: им могут потребоваться дополнительные вычисления и новый подбор гиперпараметров.
Сжимать можно второй момент, но не направление шага
Диагностика состояний Adam показала, что соседние значения второго момента похожи друг на друга сильнее, чем значения первого. Его эффективный ранг ниже: большая часть информации сосредоточена в меньшем числе направлений. После вычисления обратного квадратного корня, который Adam использует для выбора масштаба шага, различия становятся ещё меньше.
Первый момент ведёт себя иначе. Он содержит положительные и отрицательные значения и переносит информацию о направлении обновления. Когда в AdamPM тем же способом объединяли элементы первого момента, качество заметно падало даже при небольших группах. Частичное исключение отдельных слоёв из такого сжатия разрыв с Adam не устранило.
Разница между состояниями видна и в самих обновлениях весов. Для большинства шагов косинусная близость обновлений AdamPM и Adam превышала 0,993. Кривые перплексии — метрики, где меньшее значение означает, что модель лучше предсказывает следующий токен, — почти совпали на протяжении обучения.
Чем крупнее группа, тем сильнее обновление отклоняется от Adam. При увеличении модели отклонение, напротив, уменьшалось. Это поддерживает идею сжатия второго момента, но не доказывает, что максимальная группа подойдёт любой архитектуре.
AdamPM стоит проверять как замену Adam без нового подбора режима
На LLaMA 1,3B с группой из восьми элементов сумма памяти параметров и состояний оптимизатора сократилась на 29%, а перплексия составила 14,73 против 14,83 у Adam. Разница по качеству мала и оказалась в пользу AdamPM, хотя считать её самостоятельным улучшением модели по одному запуску не стоит.
Экономия относится к статической памяти параметров и состояний в BF16, а не ко всему пику памяти GPU. Активации, градиенты, временные буферы и особенности распределённого обучения могут изменить итог на конкретной инфраструктуре.
Проверки охватили предварительное обучение LLaMA на английской части C4, полную донастройку RoBERTa-base на GLUE и мультимодальную донастройку TinyLLaVA. В последних двух сценариях AdamPM сохранил среднее качество Adam при меньшем объёме состояний. Темп обучения для Adam переносили без дополнительного поиска.
Для команды, которая уже обучает модели с Adam, практический план меняется умеренно: AdamPM можно добавить как совместимый вариант оптимизатора и проверить на существующей конфигурации. Наибольший смысл это имеет, когда состояния Adam заметно ограничивают размер модели, пакета или контекста, а переход на другой оптимизатор потребовал бы новой серии запусков.
Размер группы лучше оставить параметром эксперимента. Работа показывает устойчивый компромисс на нескольких архитектурах, но крупнейшая проверенная модель остаётся меньше масштабов, на которых обычно особенно остро не хватает памяти. Перенос результата на более крупное предварительное обучение пока следует подтверждать собственным прогоном.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
