LoRA

Задача

Полное дообучение модели на 7 миллиардов параметров требует держать в памяти не только веса. Adam хранит два момента, плюс сами градиенты — то есть четыре массива размера модели вместо одного. В fp32 это порядка 100 гигабайт под то, что содержательно является небольшой поправкой.

Отсюда вопрос: нельзя ли ограничить поправку так, чтобы она была маленькой по числу параметров, оставаясь при этом полноразмерной по действию?

$$W' = \htmlData{k=frozen}{W} + \htmlData{k=scale}{\frac{\alpha}{r}}\,\htmlData{k=update}{BA}, \qquad B \in \mathbb{R}^{m\times r},\; A \in \mathbb{R}^{r\times n}$$

Поправка имеет ранг не выше $r$ и содержит $r(n+m)$ чисел вместо $nm$. Всё остальное в этом уроке — следствия этой замены.

Арифметика ранга

Экономия равна $\frac{r(n+m)}{nm}$, и у неё есть точка безубыточности: ранг, при котором две тонкие матрицы содержат столько же чисел, сколько одна большая.

$r^\ast = \frac{nm}{n+m}, \qquad \text{для квадратной } n = m: \quad r^\ast = \frac{n}{2}$

Для матрицы $4096 \times 4096$ это $2048$. Практические ранги — от $4$ до $64$, то есть в тридцать–пятьсот раз меньше точки безубыточности. Пространства тут столько, что выбор между $r = 8$ и $r = 16$ не про память вовсе.

проекции внимания 4096×4096, ×128MLP 4096×11008, ×96малая модель 768×768, ×48

ранг r

8

всего в W

2147.48M

обучаемых

8.39M

доля

0.39%

безубыточный ранг

2048

Adam: LoRA / полное

96 / 24576 МБ

Обратите внимание на последний readout — вот где настоящая экономия. Обучаемый параметр стоит не одно место, а три: сам параметр, первый момент, второй момент. Для проекций внимания при $r = 8$ это $96$ МБ против $24$ ГБ. Замораживание $W$ убирает три массива, а не один.

Почему малый ранг вообще работает

Здесь нужна честность: теоремы нет. Есть эмпирическое наблюдение, что поправка, нужная для адаптации к задаче, имеет быстро убывающий спектр, и есть теорема Эккарта–Янга из блока 1 (урок 090), которая говорит, что если спектр убывает быстро, то усечение по рангу — наилучшее из возможных приближений.

Посмотрите на спектр: если сингулярные значения падают быстро, ранг $r$ забирает почти всю норму; если равномерны — не забирает ничего.

быстро убывающийплоский

поправка ΔW

ранга r · k = 2

ранг r

2

спектр сингулярных значений

ошибка по Фробениусу

0.378

хранимых чисел

66 / 256

Формулировка, которую стоит держать в голове: LoRA — это ставка на то, что поправка низкоранговая, а не факт про нейросети. Ставка обычно выигрывает, и известно, где она проигрывает: на задачах, требующих новых знаний, а не новой формы ответа, малого ранга не хватает, и это видно по тому, что качество упирается в потолок при росте $r$ вместо роста.

Три детали, без которых не работает

Инициализация $B = 0$. Матрица $A$ заполняется случайно, а $B$ — нулями, поэтому в начале $BA = 0$ и модель в точности равна предобученной. Если бы обе были случайными, обучение начиналось бы с испорченной модели, и первые шаги уходили бы на восстановление.

Здесь стоит быть точным, потому что «одна из двух — нулями» звучит произвольно. Выпишем градиенты: при $y = Wx + BAx$ и $\delta = \partial L/\partial y$

$\frac{\partial L}{\partial B} = \delta (Ax)^\top, \qquad \frac{\partial L}{\partial A} = B^\top \delta x^\top$

Отсюда видно, что нулём обязана быть ровно одна из матриц. Обе нулями — мёртвая точка: оба градиента обращаются в нуль и остаются нулями навсегда. Любая одна нулевая — работает: при $B = 0$ на первом шаге двигается $B$, а со второго и $A$. Зеркальный вариант ($A = 0$, $B$ случайно) тоже даёт $\Delta W = 0$ и тоже трогается с места, так что выбор между ними — не математика, а эмпирика: с $A$ случайным и $B = 0$ на практике проходят большие learning rate.

Масштаб $\alpha/r$. Без него величина поправки росла бы вместе с рангом, и learning rate пришлось бы подбирать заново под каждый $r$. С ним $r$ и $\alpha$ разделены: ранг отвечает за выразительность, $\alpha$ — за силу.

Слияние. После обучения $W' = W + \frac{\alpha}{r}BA$ считается один раз, и на инференсе никакой поправки нет вовсе — ни лишних слоёв, ни лишней задержки. Это главное отличие от адаптеров, которые добавляют в сеть новые блоки и потому платят на каждом прогоне.

Что LoRA не экономит

Тут распространено завышенное ожидание. Обратный проход обязан пройти через всю сеть: чтобы получить градиент по $A$ в первом слое, нужно протащить производную через все последующие слои, замороженные они или нет.

что	полное дообучение	LoRA
параметры	все	$0.1{-}1\%$
градиенты	все	только по $A$, $B$
состояние Adam	$2\times$ все	$2\times$ по $A$, $B$
активации для backward	все	все
проход назад через сеть	целиком	целиком

Экономия — в трёх верхних строках, и она огромна. В двух нижних — нулевая. Поэтому LoRA снимает ограничение по памяти под оптимизатор, но не делает шаг существенно быстрее: время по-прежнему уходит в прямой и обратный проходы.

Отсюда и QLoRA: если веса всё равно заморожены, их можно держать в 4 битах — они не обновляются, поэтому точность их представления нужна только для прямого прохода. Обучаемая часть при этом остаётся в fp16. Комбинация из двух независимых наблюдений, каждое из которых само по себе простое.

Итог

• Замена одна: $\Delta W = BA$ ранга $r$, и всё остальное — арифметика.
• Безубыточный ранг $\frac{nm}{n+m}$ огромен по сравнению с используемыми, поэтому вопрос выбора $r$ — про качество, а не про память.
• Экономятся градиенты и состояние оптимизатора; активации и время проходов — нет.
• $B = 0$ на старте, масштаб $\alpha/r$, слияние после обучения — три детали, каждая из которых решает конкретную проблему.