Параметрически эффективное дообучение

Полное дообучение в предыдущей заметке дало 0.890 на SST-2 и заплатило за это градиентами и состояниями оптимизатора для всех 124 млн параметров и полной копией модели на каждую задачу. Линейное зондирование не заплатило почти ничего и дало 0.737. Вопрос в том, можно ли обучать немного параметров и всё же получить число полного дообучения. Можно, и главный метод опирается на одно наблюдение о том, что на самом деле меняет дообучение.

Обозначения. Предобученная матрица весов $W \in \mathbb R^{d_{\text{out}} \times d_{\text{in}}}$ (заморожена); её дообученная версия $W + \Delta W$; вход слоя $x$ и выход $h$. $r$ — ранг; $\sigma$ — сигмоида.

Отправная точка: всё или ничего#

Пока есть два режима. Обучать всё ($\Delta W$ свободна, $d_{\text{out}} d_{\text{in}}$ параметров на матрицу, сотни матриц) или обучать только выходной слой ($\Delta W = 0$ везде). Первое дорого, второе недообучается. Методы ниже лежат между ними: они ограничивают форму $\Delta W$ или добавляют небольшой модуль, параметры которого — единственные обучаемые.

Рисунок

Где каждый метод PEFT добавляет обучаемые параметры

Адаптеры#

Houlsby et al. (2019) вставляют небольшой модуль с узким местом после каждого подслоя каждого блока: проецируют $d$-мерную активацию вниз до $m \ll d$, применяют нелинейность, проецируют обратно вверх и прибавляют результат к активации (остаточная связь). Обучаются только веса адаптеров ($2dm$ на адаптер плюс смещения) и нормализации слоёв. При $d = 768$ и $m = 64$ в адаптере около 100 тыс. параметров, и 12-блочная модель с двумя адаптерами на блок получает 2.4 млн — 2% BERT-base — при точности в пределах пункта от полного дообучения на задачах GLUE из статьи. Цена — при инференсе: адаптеры — дополнительные слои, которые работают на каждом прямом проходе и не могут быть убраны.

LoRA: обновление низкого ранга#

Hu et al. (2022) построили низкоранговую адаптацию (LoRA) на наблюдении — эмпирическом и сделанном раньше Aghajanyan, Zettlemoyer and Gupta (2021), — что обновление, которое нужно предобученной модели для прикладной задачи, имеет низкий внутренний ранг (intrinsic rank): бо́льшая часть направлений в $\Delta W$ не важна. Слой вычисляет

$$h = W x + \Delta W x = W x + B A x = W x + B\,(A x),$$

справа налево: $Ax$ — это $r$ чисел, $B(Ax)$ — снова $d_{\text{out}}$ измерений. $A$ инициализируется маленькими гауссовскими значениями, а $B$ — нулями, так что на первом шаге $\Delta W = 0$, и обучение начинается ровно с предобученной модели. Множитель $\alpha / r$ масштабирует поправку, чтобы изменение $r$ не меняло эффективный шаг обучения.

Рисунок

LoRA: замороженная матрица весов плюс обучаемая поправка низкого ранга

Разобранный пример#

$d_{\text{in}} = d_{\text{out}} = 4$, $r = 1$. Полное обновление — 16 чисел; LoRA обучает $A = (1, 0, -1, 0)$ (4 числа) и $B = (2, 0, 0, 1)^\top$ (4 числа). Их произведение

$$BA = \begin{pmatrix} 2 \\ 0 \\ 0 \\ 1 \end{pmatrix}\begin{pmatrix} 1 & 0 & -1 & 0 \end{pmatrix} = \begin{pmatrix} 2 & 0 & -2 & 0 \\ 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 \\ 1 & 0 & -1 & 0 \end{pmatrix},$$

— матрица $4 \times 4$ ранга 1: каждая строка кратна $A$. Для $x = (1, 1, 0, 1)$ получаем $Ax = 1$ и $BAx = (2, 0, 0, 1)$: поправка добавляет направление $B$, масштабированное тем, насколько $x$ сонаправлен с $A$. Восемь обученных чисел задают матрицу из 16 элементов ценой того, что матрица может указывать только в одном направлении. С $r$ направлениями — в $r$.

Число параметров#

На матрицу: $r(d_{\text{in}} + d_{\text{out}})$ против $d_{\text{in}} d_{\text{out}}$. Для слитой проекции внимания GPT-2 ($768 \to 2\,304$, дающей $Q, K, V$ сразу) при $r = 8$: $8 \times (768 + 2\,304) = 24\,576$ на блок, 294 912 на 12 блоков, плюс выходной слой на 1536 параметров: 296 448 обучаемых параметров, 0.24% модели. Исходная статья применяет LoRA только к проекциям запросов и значений; здесь она применяется к слитой проекции GPT-2, которая покрывает запросы, ключи и значения, потому что GPT-2 вычисляет их одной матрицей.

Результат#

Режим	Обучается	Доля верных (валидация)	Потери (валидация)
Полное дообучение	124 млн (100%)	0.890	0.277
Линейное зондирование	1536 (0.001%)	0.737	0.541
LoRA, $r = 8$ на внимании	296 тыс. (0.24%)	0.899	0.271
Настройка промпта, 12 токенов	10.8 тыс. (0.009%)	0.663	0.656
LP-FT	124 млн	0.900	0.258

LoRA не уступает полному дообучению (разница в пункт — шум на 872 примерах), обучая в 400 раз меньше параметров, с шагом обучения $10^{-4}$ (новые матрицы начинают с нуля, поэтому осторожность с маленьким шагом из полного дообучения не нужна) и без состояний оптимизатора для замороженного тела. Адаптер задачи — файл из 296 тыс. чисел, который можно подключать к одной общей базовой модели и отключать.

Настройка промпта#

Lester, Al-Rfou and Constant (2021) не обучают ничего внутри модели. К эмбеддингам каждого входа приписывается последовательность из $k$ виртуальных токенов — свободных векторов размерности $d$; модель заморожена; обучаются только эти $k \times d$ чисел (и здесь выходной слой). Метод опирается на наблюдение, что большими языковыми моделями можно управлять промптом, и заменяет написанные вручную промпты выученными. Инициализация виртуальных токенов эмбеддингами естественного предложения (Predict if sentiment of this review is positive or negative, 11 токенов, дополненных до 12) даёт оптимизатору разумную отправную точку.

Результат: 0.663 при 10 752 обучаемых параметрах — ниже линейного зондирования. Эксперименты статьи поставлены на моделях энкодер–декодер T5 на SuperGLUE, и там настройка промпта (prompt tuning) догоняет полное дообучение только около самого большого из опробованных размеров, примерно 11 млрд параметров, и сильно отстаёт на меньших (Lester et al., 2021, рисунок 1). Эта кривая — свидетельство о T5, а не правило для любой архитектуры, но она указывает в ту же сторону, что и этот прогон: GPT-2 на 124 млн параметров далеко от размера, на котором разрыв закрылся. Сами мягкие промпты не привязаны к одной архитектуре. Статья приписывает их ко входу модели энкодер–декодер, и тот же механизм применим к энкодеру и к модели из одного декодера; насколько хорошо он работает, зависит от модели и её размера и измеряется для каждой постановки. Его привлекательность — в размере: задача — это несколько тысяч чисел.

Более сильный родственник — настройка префикса (prefix tuning; Li and Liang, 2021): выученные векторы подставляются как дополнительные ключи и значения в каждый слой внимания, а не только на вход.

Дообучение квантизованной модели: QLoRA#

LoRA держит замороженное тело в памяти, и для модели на 7 млрд это 14 ГБ в 16 битах. Dettmers et al. (2023) хранят замороженное тело в 4 битах (формат NF4 из заметки о квантизации, 3.5 ГБ для 7 млрд), деквантуют каждую матрицу на лету для прямого прохода, держат матрицы LoRA в 16 битах и применяют вторую квантизацию к самим константам квантизации. Градиенты проходят через деквантованные веса только в $A$ и $B$. Тот же рецепт работает и с моделью, квантованной методом GPTQ: для 4-битной Llama-2-7B LoRA с $r = 8$ на четырёх проекциях внимания добавляет 8.4 млн обучаемых параметров, и десять шагов обучения на наборе цитат выполняются на одном потребительском GPU. В экспериментах статьи качество дообучения QLoRA близко к 16-битной LoRA; механизм, который это позволяет, в том, что замороженным весам никогда не нужны градиенты, поэтому формат их хранения может быть любым, лишь бы его мог прочитать прямой проход.

Компромиссы#

Метод	Обучается	Цена при инференсе	Где применим	Качество (на этой задаче)
Полное дообучение	100%	нет	везде	0.890
Адаптеры	~2%	дополнительные слои	везде	≈ полное (литература)
LoRA	0.1–1%	нет после слияния	линейные слои	0.899
Настройка промпта	< 0.01%	$k$ дополнительных токенов на вход	любые; конкурентна в большом масштабе (T5 в статье)	0.663
QLoRA	LoRA на 4-битном теле	4-битное тело	большие декодеры	≈ LoRA (литература)

Практические решения: $r$ от 4 до 64 (статья обнаружила, что для GPT-3 на её задачах достаточно значений от $r = 1$ до $8$, а дальше качество не растёт); применять сначала к проекциям внимания, а при наличии бюджета — и к матрицам полносвязных слоёв; шаг обучения около $10^{-4}$ – $3 \times 10^{-4}$; дропаут на пути LoRA.

Что изменилось#

	Полное дообучение	LoRA
$\Delta W$	Свободна	Ранг $\le r$
Параметров на задачу	Все	$r(d_{\text{in}} + d_{\text{out}})$ на матрицу
Память оптимизатора	2 состояния на вес	2 состояния на вес LoRA
Артефакт задачи	Полная модель	Файл адаптера
Инференс	Тот же	Тот же, после слияния
Шаг обучения	Маленький	Больше (новые веса)

Ограничения#

• Ранг — это ставка. Задачам, которым могут понадобиться большие изменения высокого ранга (новый язык, долгое продолжение предобучения на новой области), малый $r$ может подходить хуже. Это компромисс ёмкости и оптимизации, который измеряют в каждом случае, а не правило, по которому LoRA не может выучить содержание области.
• Замороженное тело — ограниченная ёмкость. Всё новое хранится в $\Delta W$. Это обновление обучается и меняет выходы модели, в том числе факты; малый $r$ ограничивает, насколько сильно оно может их изменить. Адаптацию к задаче от продолжения предобучения отличают в основном цель и данные (размеченные примеры задачи против большого объёма сырого текста), и то и другое можно проводить низкоранговыми обновлениями.
• Зависимость от масштаба. Полезность настройки промпта зависит от размера модели; гиперпараметры LoRA ($r$, целевые модули, $\alpha$) зависят от модели и задачи и подбираются эмпирически.
• Адаптеры добавляют задержку, LoRA — нет; QLoRA обменивает вычисления на деквантование на память.
• Шум. Сравнения выше опираются на один запуск каждого метода на 872 примерах; разница в пункт — не вывод.

Что нужно запомнить#

• LoRA: $h = Wx + BAx$, где $W$ заморожена, $A$ случайна, а $B = 0$ в начале; $r(d_{\text{in}} + d_{\text{out}})$ обучаемых параметров на матрицу; для инференса $BA$ вливается в $W$.
• На SST-2 с GPT-2 small LoRA ($r = 8$, внимание) даёт 0.899 против 0.890 у полного дообучения при обучении 0.24% параметров.
• Адаптеры добавляют небольшие слои с узким местом (лишняя цена инференса); настройка промпта выучивает входные векторы, в экспериментах Lester et al. на T5 догнала полное дообучение только около 11 млрд параметров, а на GPT-2 small здесь дала 0.663.
• QLoRA хранит замороженное тело в 4 битах и обучает поверх LoRA, что позволяет дообучать модели на 7 млрд параметров на одном GPU.
• Ранг — гипотеза о том, насколько задача должна изменить модель. Для адаптации к задаче обычно хватает малого ранга; большим изменениям, например новой области, может понадобиться больший $r$, больше целевых модулей или полное обучение, и выбор проверяют на валидационных данных.

Ноутбук#

week10_PEFT.ipynb. Номера ячеек считают все ячейки, включая текстовые, начиная с 1. Приведённые числа взяты из прогона, сохранённого в курсе; у ячеек, изменённых с тех пор под текущие версии библиотек, сохранённого выхода нет.

• Ячейки 35–39 настраивают LoraConfig(task_type=SEQ_CLS, r=8, lora_dropout=0.1, target_modules=['c_attn']), где c_attn — слитая проекция запросов, ключей и значений GPT-2, и оборачивают модель в get_peft_model; ячейка 39 насчитывает 296 448 обучаемых параметров. В сохранённом выходе ячейки 38 нет строки print_trainable_parameters; вместе с матрицами LoRA в модели 124 738 560 параметров, так что обучаемая доля — 0.24%.
• Ячейки 40–43 обучают с AdamW при $10^{-4}$ 1000 шагов и приводят функцию потерь на валидации 0.2708 и долю правильных ответов 0.8991.
• Ячейки 45–49 настраивают PromptTuningConfig с 12 виртуальными токенами, инициализированными приведённым выше предложением (num_virtual_tokens=12, token_dim=768), и печатают 10 752 обучаемых параметра, а ячейки 52–54 обучают и приводят функцию потерь 0.6560 и долю правильных ответов 0.6628.

week12_pt1_Distillation_and_GPTQ_Quantization.ipynb, ячейки 41–51: 4-битная GPTQ-модель Llama-2-7B подготавливается prepare_model_for_kbit_training, оборачивается LoraConfig(r=8, lora_alpha=32, target_modules=["k_proj","o_proj","q_proj","v_proj"], lora_dropout=0.05), что печатает 8 388 608 обучаемых параметров, и обучается 10 шагов через Trainer; напечатанное там число «all params» (270 млн) считает упакованные 4-битные веса как отдельные параметры int32 и не равно настоящему размеру модели.

week12_pt0_Torch_and_ONNX_Quantization_+_PEFT.ipynb, ячейки 71–76: конфигурация QLoRA (load_in_4bit, nf4, двойная квантизация, вычисления в bfloat16) и LoraConfig(r=16, target_modules="all-linear") для базовой Mistral-7B, то есть LoRA ставится и на матрицы полносвязных слоёв, а не только на проекции внимания; эти ячейки — шаблон и до конца не запускаются.

Дальше#

До сих пор адаптация означала размеченную задачу: тональность, тему. Модели, с которыми люди действительно взаимодействуют, адаптированы к чему-то более широкому — к поведению «отвечать на инструкцию», и следующая заметка описывает этот процесс: обучение с учителем на парах «инструкция — ответ», затем оптимизацию по предпочтениям людей, причём LoRA и QLoRA — обычный способ себе это позволить.

Литература#

• E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, W. Chen. LoRA: Low-Rank Adaptation of Large Language Models. ICLR, 2022.
• A. Aghajanyan, L. Zettlemoyer, S. Gupta. Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning. ACL, 2021.
• N. Houlsby et al. Parameter-Efficient Transfer Learning for NLP. ICML, 2019. Адаптеры.
• B. Lester, R. Al-Rfou, N. Constant. The Power of Scale for Parameter-Efficient Prompt Tuning. EMNLP, 2021.
• X. L. Li, P. Liang. Prefix-Tuning: Optimizing Continuous Prompts for Generation. ACL, 2021.
• T. Dettmers, A. Pagnoni, A. Holtzman, L. Zettlemoyer. QLoRA: Efficient Finetuning of Quantized LLMs. NeurIPS, 2023.
• Hugging Face. PEFT documentation. Библиотека, используемая в ноутбуках.