Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Матричные разложения в LoRA и слоях внимания можно обучать так, чтобы один множитель не рос неконтролируемо при неизменном результате. При полном предобучении GPT2 новый оптимизатор дал сопоставимую ошибку при той же пиковой памяти; препринт не рецензирован, а все числа получили сами авторы. Stiefel-AdamW позволяет проверить эту замену на существующем факторизационном блоке, не меняя всю архитектуру или схему обучения.
Почему одинаковая матрица создаёт разные условия для AdamW
Линейный факторизационный блок хранит матрицу W как произведение B·A. Такая конструкция встречается в адаптерах LoRA, сжатых слоях и произведениях query-key в механизме внимания.
Одно и то же W можно получить из множества пар A и B. Один фактор способен уменьшаться, а другой — расти, пока их произведение остаётся прежним. Функция потерь при этом почти не меняется, но сами параметры и градиенты принимают неудобный для численных расчётов масштаб.
AdamW обрабатывает координаты факторов по отдельности и не учитывает эту неоднозначность. Несбалансированные факторы ограничивают допустимую скорость обучения и могут сделать процесс нестабильным, хотя эффективная матрица модели остаётся ограниченной.
Stiefel-AdamW требует, чтобы строки A оставались ортонормированными: каждая имеет единичную длину и перпендикулярна остальным. Такое множество матриц называют многообразием Штифеля. Ограничение убирает преобразования, при которых один фактор может неограниченно расти за счёт другого, но оставляет компактное семейство ортогональных замен координат.
Геометрическая поправка применяется только перед обновлением
Для обоих факторов оптимизатор накапливает среднее градиента и среднее его квадратов в обычном евклидовом пространстве. Поэтому сохраняется покоординатное масштабирование, ради которого AdamW используют на практике.
Фактор B получает стандартный шаг AdamW. Для A готовое адаптивное направление сначала проецируется на допустимое направление движения, а затем матрицу возвращают на множество ортонормированных матриц. Последнюю операцию называют ретракцией; работа рассматривает варианты на основе преобразования Cayley, QR-разложения, полярного разложения и итерации Newton–Schulz.
Геометрия входит только в последние две операции. Не требуется хранить отдельную модель или отказываться от состояний AdamW, а разные способы ретракции дали близкие результаты. Авторы по умолчанию использовали приближённое преобразование Cayley.
На полном предобучении GPT2 тестовая ошибка составила 3,236 против 3,260 у AdamW — разница находится в пределах разброса. Пиковое потребление памяти совпало и составило 13,8 ГБ. В предобучении Qwen2 с адаптерами ошибка снизилась до 2,59 против 2,66.
При дообучении Mistral 7B на GLUE метод показал лучший средний результат среди рассмотренных оптимизаторов и лидировал на MNLI, MRPC, STS-B и WNLI. На ViT для CIFAR-10 он сходился со скоростью, сопоставимой с AdamW и Scaled AdamW, но сохранял устойчивость при более агрессивной скорости обучения.
Проверки охватывают LoRA-дообучение GPT2, ViT и Mistral 7B, полное предобучение GPT2 на OpenWebText и предобучение Qwen2 с адаптерами на WikiText-103. Основными соперниками выступали AdamW, Scaled AdamW и GeoLoRA, но условия и набор сравнений различались между задачами.
Планы менять стоит на уровне оптимизатора, а не архитектуры
Работа не требует пересматривать выбор LoRA или заменять трансформер. Она предлагает другой способ обновлять уже существующую пару факторов, если они перемножаются напрямую и между ними нет нелинейной функции.
Для команды это означает локальное изменение: выделить связанные параметры A и B, ортонормировать ограничиваемый фактор при инициализации и добавить проекцию с ретракцией в оптимизатор. Остальные веса модели продолжает обновлять обычный AdamW. Наиболее прямой сценарий проверки — нестабильное LoRA-дообучение, где сейчас приходится снижать скорость обучения или вручную следить за нормами факторов.
Результаты пока не обосновывают обязательную миграцию всех проектов. Теоретическая часть доказывает ограниченность градиентов и свойства функции сожаления при стандартных предположениях выпуклой оптимизации, но обучение нейросетей в экспериментах остаётся невыпуклым. Кроме того, метод устраняет опасное масштабирование факторов, но сохраняет неоднозначность при ортогональной смене координат.
Практический план — реализовать Stiefel-AdamW как отдельную группу параметров и сравнить его с текущим AdamW на целевой задаче при одинаковых настройках. Если проблема связана с ростом норм факторов или чувствительностью к скорости обучения, работа даёт конкретный способ стабилизации без дополнительной памяти в показанном опыте. Для моделей без обучаемых факторизационных блоков она планов не меняет.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
