Новый подход к обучению языковых моделей снизил затраты памяти без потери качества

Новый фреймворк FRUGAL разработан для обучения больших языковых моделей. Метод позволяет значительно сократить объем памяти, необходимый для хранения исторических статистик оптимизатора, при этом сохраняя качество обучения. Результаты исследования опубликованы на ICML 2025 — одной из ведущих конференций в области машинного обучения. Текст статьи можно посмотреть также на arXiv.

«Главная идея FRUGAL и отличие от предыдущих подходов в том, что остаточная часть градиента после проекции все еще хранит в себе полезную информацию, которую все еще можно использовать для обучения модели. Наш подход позволяет сильно демократизировать обучение больших моделей в условиях дефицита графических ускорителей», — объяснил Александр Безносиков, научный руководитель BRAIn Lab, директор Центра агентных систем Института искусственного интеллекта МФТИ, заведующий лабораторией проблем федеративного обучения ИСП РАН.

Современные языковые модели, такие как GPT и LLaMA, содержат миллиарды параметров. При их обучении значительная часть памяти видеокарты уходит не на сами веса модели, а на служебную информацию в оптимизаторе. Например, популярный алгоритм Adam хранит для каждого параметра две дополнительные величины: скользящие средние градиента и его квадрата. Для модели с 8 миллиардами параметров это означает дополнительные 64 гигабайта памяти — вместе с самой моделью получается больше, чем вмещают многие профессиональные серверные видеокарты.

Существующие методы экономии памяти, такие как LoRA и GaLore, решают проблему за счет работы в пространстве пониженной размерности. Однако при этом теряется часть информации из градиента, что может ухудшить результат обучения.

Авторы предложили принципиально иной подход: разделить пространство параметров на две части. Для первой части используется продвинутый оптимизатор (например, Adam), требующий хранения статистик. Для второй — простой метод без статистик, такой как signSGD, который не требует дополнительной памяти вовсе.

Ключевое наблюдение в том, что не все параметры модели одинаково «требовательны» к выбору оптимизатора. Эксперименты показали неожиданный результат: только выходной слой критически нуждается в Adam, тогда как остальные компоненты, включая слои нормализации и эмбеддинги, можно обучать простым signSGD практически без потери качества.

«Знаете, в некотором смысле это было очень похоже на то, как в OpenAI масштабировали обучение моделей для игры в Dota 2 — мы также были абсолютно уверены, что в какой-то момент с понижением размерности все сломается и результаты резко ухудшатся. Но мы продолжали и продолжали понижать размерность, а результаты практически не менялись. В итоге мы дошли до размерности 0, что фактически означает полное обучение простым signSGD. Это было очень неожиданно и довольно сильно противоречит общепринятой в области интуиции. Мы уже планируем следующий проект на основе этого результата», — рассказал Филипп Змушко, сотрудник BRAIn Lab МФТИ.

На задачах предобучения моделей семейства LLaMA (от 60 миллионов до трех миллиардов параметров) FRUGAL стабильно превосходит конкурирующие подходы GaLore и BAdam при одинаковом бюджете памяти. При этом метод практически не отстает от полнорангового Adam — стандартного алгоритма, который требует в несколько раз больше памяти для хранения состояния оптимизатора.

Авторы также доказали теоретические гарантии сходимости алгоритма, что подтверждает его надежность.

В работе приняли участие ученые из Центра агентных систем Института искусственного интеллекта МФТИ, «Яндекса» и Университета искусственного интеллекта имени Мохамеда бин Заида (ОАЭ).