Тонкая настройка языковых моделей с использованием только прямого прохода
Fine-Tuning Language Models with Just Forward Passes
May 27, 2023
Авторы: Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D. Lee, Danqi Chen, Sanjeev Arora
cs.AI
Аннотация
Тонкая настройка языковых моделей (LM) показала успех в решении разнообразных задач, однако с увеличением размера моделей обратное распространение требует непомерно большого объема памяти. Методы нулевого порядка (ZO) теоретически могут оценивать градиенты, используя всего два прямых прохода, но считается, что они катастрофически медленны для оптимизации крупных моделей. В данной работе мы предлагаем оптимизатор нулевого порядка с эффективным использованием памяти (MeZO), адаптируя классический метод ZO-SGD для работы на месте, что позволяет выполнять тонкую настройку LM с тем же объемом памяти, что и при выводе. Например, с использованием одной видеокарты A100 80GB MeZO может обучать модель с 30 миллиардами параметров, тогда как тонкая настройка с обратным распространением позволяет обучать только LM с 2,7 миллиардами параметров при том же бюджете. Мы провели всесторонние эксперименты с различными типами моделей (маскированные и авторегрессивные LM), масштабами моделей (до 66 миллиардов параметров) и задачами (классификация, множественный выбор и генерация). Наши результаты показывают, что (1) MeZO значительно превосходит обучение в контексте и линейное зондирование; (2) MeZO достигает сопоставимой производительности с тонкой настройкой с обратным распространением в различных задачах, при этом сокращая использование памяти до 12 раз; (3) MeZO совместим как с полной настройкой параметров, так и с методами эффективной настройки, такими как LoRA и префиксная настройка; (4) MeZO может эффективно оптимизировать недифференцируемые цели (например, максимизацию точности или F1). Мы подкрепляем наши эмпирические результаты теоретическими выводами, подчеркивая, как адекватная предварительная подготовка и подсказки задач позволяют MeZO выполнять тонкую настройку огромных моделей, несмотря на то, что классический анализ ZO предполагает обратное.
English
Fine-tuning language models (LMs) has yielded success on diverse downstream
tasks, but as LMs grow in size, backpropagation requires a prohibitively large
amount of memory. Zeroth-order (ZO) methods can in principle estimate gradients
using only two forward passes but are theorized to be catastrophically slow for
optimizing large models. In this work, we propose a memory-efficient
zerothorder optimizer (MeZO), adapting the classical ZO-SGD method to operate
in-place, thereby fine-tuning LMs with the same memory footprint as inference.
For example, with a single A100 80GB GPU, MeZO can train a 30-billion parameter
model, whereas fine-tuning with backpropagation can train only a 2.7B LM with
the same budget. We conduct comprehensive experiments across model types
(masked and autoregressive LMs), model scales (up to 66B), and downstream tasks
(classification, multiple-choice, and generation). Our results demonstrate that
(1) MeZO significantly outperforms in-context learning and linear probing; (2)
MeZO achieves comparable performance to fine-tuning with backpropagation across
multiple tasks, with up to 12x memory reduction; (3) MeZO is compatible with
both full-parameter and parameter-efficient tuning techniques such as LoRA and
prefix tuning; (4) MeZO can effectively optimize non-differentiable objectives
(e.g., maximizing accuracy or F1). We support our empirical findings with
theoretical insights, highlighting how adequate pre-training and task prompts
enable MeZO to fine-tune huge models, despite classical ZO analyses suggesting
otherwise.