Как дообучить LLM, когда память на исходе 🧐

Сегодня представляем статью Leveraging Coordinate Momentum in SignSGD and Muon: Memory-Optimized Zero-Order, написанную совместно с RnD центром T-Bank,  где мы исследуем Zero-Order (ZO) оптимизацию как эффективную альтернативу для fine-tuning больших языковых моделей, особенно в сочетании с техниками вроде LoRA

[статья]

Проблема: непомерные затраты
памяти💀

Традиционные методы оптимизации первого порядка, такие как SGD и Adam, требуют огромных вычислительных ресурсов и памяти, что становится серьезным препятствием при работе с большими моделями. В предыдущем посте мы рассказывали о методе FRUGAL, который позволял существенно сократить использование памяти за счет уменьшения объема хранимых статистик оптимизатора. Однако FRUGAL был разработан для общей задачи обучения больших языковых моделей с нуля. Когда же речь идет о дообучении (fine-tuning) уже предобученных моделей, можно применить еще более экономные подходы — оптимизацию нулевого порядка.

Главное преимущество ZO-подхода: для оценки градиента вообще не требуется вычислять обратное распространение ошибки (loss.backward()) 🤯

Вместо этого градиент приближается с помощью простых вычислений forward функции потерь при небольших сдвигах параметров:
∇loss(w) ≈ [loss(w + τe) - loss(w - τe)] / (2τ) * e,
где w – обучаемые параметры (веса модели), 
e — случайный вектор,
τ — малый параметр сглаживания

Хотя методы нулевого порядка уже активно применяются для дообучения LLM, до сих пор большинство работ ограничивались базовыми алгоритмами без дополнительных техник. В этой статье мы расширяем возможности ZO-оптимизации, вводя продвинутые методы с momentum: мы адаптируем ZO подход JAGUAR 🐆 к алгоритмам SignSGD и Muon в безградиентной постановке🔞

Мы выбрали именно эти алгоритмы, потому что они уже показали отличные результаты в задачах дообучения LLM, зачастую превосходя даже такие популярные методы, как AdamW (про SignSGD у нас уже выходил отдельный пост ранее)

Ключевые инновации 🤓

1. JAGUAR SignSGD — первый ZO алгоритм с momentum для SignSGD:
- Использует coordinate-wise обновления (по одной координате за раз);
 - Накапливает информацию из прошлых итераций через momentum β;
 - Требует всего 2d + 1 параметров (вместо 4d у ZO Adam), где d – размерность входного вектора w;
 - Сходится в невыпуклом случае (элегантная скорость сходимости – в комментариях)

2. JAGUAR Muon — ZO версия современного матричного оптимизатора Muon:
 - Работает с матричной структурой параметров;
 - Использует Newton-Schulz ортогонализацию;
 - Первый ZO Muon в литературе;
 - Сходится в невыпуклом случае (элегантная скорость сходимости – в комментариях)

Экспериментальные результаты

Тестирование на SST2, COPA, WinoGrande с моделями OPT-1.3B, RoBERTa-Large, Llama2-7B, OPT-13B показало:

✅ Метрики: Accuracy у JAGUAR методов превосходит все ZO бейзлайны и иногда даже First Order методы;
✅ Память: ~4.17 GB vs 12.2 GB у FO-SGD в точности F16 (экономия в 3 раза!);
✅ Масштабируемость: работает как на небольших, так и на моделях с 13B параметрами 

Практическое значение

Эти методы открывают возможности для:
- Адаптации больших языковых моделей на обычных рабочих станциях и даже ноутбуках;  
- Получения сопоставимых результатов с методами первого порядка, но с меньшими требованиями к ресурсам; 
- Более гибкого применения LLM в ситуациях, где классические методы невозможны или неэффективны

Так, Zero-order оптимизация из экзотической техники превращается в практический инструмент для демократизации AI

#разборстатьи

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