Можно ли обучить языковую модель, используя signSGD?
Оказывается, да — но с небольшой оговоркой.

Сегодня расскажем о совместной с MBZUAI (Mohamed bin Zayed University of Artificial Intelligence, ОАЭ) статье FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training (Ф. Змушко, А. Безносиков, М. Такач, С. Хорват)

[статья]

Введение

Современные языковые модели показывают невероятные результаты в самых разных задачах. Но этот прогресс достигается за счёт роста размера моделей, что приводит к резкому увеличению требований к памяти.

Например, для обучения модели на 8B параметров с использованием стандартного AdamW, нужно хранить:

- веса: 8B
- градиенты: 8B
- статистики оптимизатора: 8B для m (моментум) и 8B для v (предобработчик)

Итого: (8 + 8 + 8 + 8 × 4 байта = 128 GB (в формате float32)
А это значительно больше, чем доступные 80 GB на A100/H100 — и это ещё без учёта активаций!

Поэтому возникает естественное желание — сократить память, особенно на хранение статистик оптимизатора. Именно этого мы и пытается добиться через FRUGAL.

Предыдущие подходы

Наиболее популярные методы — это PEFT-подходы (например, LoRA), которые добавляют к матрице весов линейного слоя обучаемые низкоранговые матрицы: W + AB

Но:
- это работает только для fine-tuning'а (и то не всегда)
- совершенно не годится для обучения с нуля

ReLoRA решает это, вмердживая произведение Aᵢ ⋅ Bᵢ в W раз в T шагов и инициализируя новые Aᵢ₊₁, Bᵢ₊₁.

GaLore идёт дальше: она факторизует градиент G, а не веса W, и обновление идёт через proj(G).

FRUGAL — гибрид с экономией памяти

Хотя итоговое изменение весов в ReLoRA и GaLore — высокоранговое, каждое отдельное обновление — низкоранговое.
Это кажется расточительным, ведь используется лишь часть информации из градиента G.

В работе мы предлагаем через FRUGAL не выбрасывать остаток G - proj(G), а использовать его для обновлений.
Но так как цель — экономия памяти, эти обновления должны быть state-free, например signSGD или SGD.

Итоговый гибридный подход представлен в Алгоритме 1 (см. в комментариях):

1. На каждом шаге градиент G разбивается на две части: G = proj(G) + (G - proj(G))
2. proj(G) — используется для state-full обновления (например, AdamW)
3. G - proj(G) — идёт на state-free обновление (например, signSGD)

Менять можно:
- оптимизаторы для обеих частей
- тип проекции: RandK, поблочная проекция и др.

Смотри продолжение (эксперименты и выводы) в комментариях ⬇️