Итак, мы прочитали Шиншиллу, Каплана и прочие работы про scaling laws и разобрали их. Мы знаем рецепты: хотим модель на N параметров — прикинули, сколько нужно токенов и FLOPs и какой будет loss. Даже размер батча подбирать научились.
Вроде, все супер.
Но есть вопрос, на который scaling laws не отвечают: с каким learning rate обучать модель?
Learning Rate — это размер шага, которым оптимизатор двигает веса на каждой итерации. Если он слишком большой, то обучение штормит или взрывается (loss улетает в бесконечность). Если он слишком маленький, то модель ползет к минимуму слишком медленно и за бюджет токенов не доходит до хорошего качества. Где-то посередине есть оптимум, но он свой для каждой конфигурации.
И вот в чем проблема: оптимальный LR зависит от размера модели, а точнее, от ее ширины. Подобрал идеальный LR на модели в 100M параметров, и на модели 100B он уже совсем другой.
Остается только перебирать заново, что ну очень дорого: один прогон модели на 175B стоит миллионы долларов и недели GPU-времени, а перебор — это десятки-сотни таких прогонов.
Поэтому исторически LR для гигантов подбирали полу-вслепую: на основании метода чуйки, по опыту прошлых моделей и с запасом осторожности: брать поменьше, лишь бы не взорвалось, и тем самым модели недоучивали.
μP (maximal update parametrization) и μTransfer убирают это гадание на рунах из абрикосовой косточки.
Идея заключается в том, чтобы сделать дорогой перебор один раз на маленькой и дешевой модели, а потом перенести найденный LR на большую модель по некому правилу масштабирования, которое задается в методе, без всяких эмпирических подгонок.
01Что такое μP
Кроме гипер параметров самой архитектуры (сколько у модели слоев, какая у них ширина, какой тип attention используется, какое количество нейронов) у модели есть еще два гипер параметра, про которые редко говорят.
- С каких чисел стартуют веса. Перед обучением веса заполняют случайным шумом. И вариация тут заключается в том, на сколько крупный этот шум или, другими словами, какой у него разброс.
- Насколько большой шаг делает каждый слой при обучении. Это learning rate, и не обязательно один на всех: у разных слоев он может быть разным.
Вот эти два набора правил: каким шумом инициализировать и какой LR у каждого слоя, — называются параметризацией.
Это не веса и не архитектура, а правила настройки вокруг них. Одну и ту же модель можно стартовать по разным таким правилам, и обучаться она будет от этого сильно по-разному.
Standard Parametrization / SP — это дефолтные правила
SP (standard parametrization) — это то, что получается по умолчанию, ничего специально не делая:
- каждый слой инициализируется шумом с разбросом ~1/√fan_in (это Kaiming/LeCun init, зашитый в PyTorch);
- один и тот же learning rate на все слои;
- никакого особого обращения с первым / последним слоем или с attention.
Просто и работает.
Но есть скрытая беда: при этих правилах поведение модели меняется, когда меняешь ширину.
Сделал слой шире, и все суммы внутри стали крупнее, шаг обновления де-факто стал больше, и тот LR, что был идеален на узкой модели, на широкой уже мимо.
μP — это те же правила, но с поправкой на ширину
μP (maximal update parametrization) — это та же модель (архитектура, число параметров, слои, все это идентично), но правила инициализации и LR прописаны уже с учетом ширины, чтобы дрейф, описанный выше, не происходил:
- скрытые слои инициализируются так же (~1/√fan_in), но их learning rate уменьшается по мере роста ширины (для Adam — делится на то, во сколько раз модель шире базовой);
- первый (embedding) и последний (readout) слои настраиваются по отдельным правилам;
- attention-логиты делятся на d_head, а не на √d_head.
Смысл всех этих поправок один: сделать так, чтобы один шаг обучения менял модель на одну и ту же величину при любой ширине. Тогда узкая и широкая версии ведут себя одинаково, и правильный LR — это одно и то же число хоть на 256, хоть на 4096.
Отсюда и название maximal update: масштабы подобраны так, чтобы каждый слой на каждом шаге обновлялся по максимуму осмысленно — не затухая (слишком мелкий шаг → слой не учится) и не взрываясь (слишком крупный → разносит), — и так на всех ширинах.
В чём отличие от SP, в одной мысли
Разница не в том, что μP делает какую-то другую модель. Разница в том, что происходит, когда ты меняешь ширину своей модели:
- в SP правила фиксированы, а поведение модели плывет, оптимальный LR уезжает, и в итоге нам надо перебирать заново на большой модели;
- μP: правила подстраиваются под ширину, а поведение остаётся стабильным, оптимальный LR стоит на месте, и мы можем переносить подобранный LR с маленькой модели на большую.
То есть μP — это умная система координат для настройки, в которой изменение ширины перестает сбивать подобранные гиперпараметры.
Как это выглядит на глаз (before / after)
Возьмём сеть embedding → скрытый слой → скрытый слой → readout и расширим скрытую ширину с 256 до 4096 (в 16 раз):
- SP: во всех слоях init ~1/√fan_in, LR = η. Это происходит одинаково и до, и после расширения. Правила не изменились, но модель на 4096 ведет себя иначе, чем на 256, и старый η уже не оптимален.
- μP: на embedding подбираем свой масштаб; скрытые слои делаем такой же init, но LR = η/16; на readout будет своя, более сильная, инициализация и выход, деленный на 16; на слоях attention будет деление на d_head. Правила подстроились под ×16: модель на 4096 с точки зрения подбираемых гипер параметров и их использования при обучении, ведет себя также, как на 256, и тот же η остается оптимальным.
Ключевой график
Самый убедительный аргумент во всей теме: loss как функция learning rate, по одной кривой на каждую ширину модели.
- Без μP (SP): у каждой кривой минимум в своем месте. Оптимальный LR ползет вправо/влево с ростом ширины. Кривые разъезжаются.
- Под μP: минимумы всех кривых стоят друг над другом. Оптимальный LR один и тот же при любой ширине.
Почему LR вообще изменяется с шириной
Представьте один нейрон внутри слоя. Он считает взвешенную сумму своих входов:
\[ y = w_1 x_1 + w_2 x_2 + \ldots + w_d x_d \]\(d\) — это ширина слоя (размерность d_model). В сумме ровно \(d\) слагаемых.
Теперь сделали слой вдвое шире: \(d\) стало \(2d\), в сумме вдвое больше членов. Если масштабы весов \(w_i\) и входов \(x_i\) не подстроить, эта сумма в среднем растет с шириной — просто потому что складывается больше чисел. Активации раздуваются. Градиенты раздуваются. Обновления весов раздуваются.
И вот learning rate, который был ок на небольшой модели (где сумма была меньше), на широкой оказывается фактически слишком большим, потому что шаг обновления уже сам по себе стал крупнее. Обучение теряет стабильность.
Это не баг и не ошибка реализации, а прямое следствие того, как устроена линейная алгебра внутри трансформера. И μP чинит это в корне: подбирает масштабы инициализации и LR так, чтобы величина обновлений оставалась порядка \(O(1)\) при любой ширине. Как только внутренний масштаб зафиксирован, LR перестаёт зависеть от ширины. Он становится константой, которую можно найти один раз.
Два правила на практике
Вся параметризация μP сводится к двум масштабированиям:
1. Инициализация. Разброс (стандартное отклонение) весов скрытых слоев масштабируется как \(1 / \sqrt{\text{fan\_in}}\).
Fan_in — это количество входов в слой, то есть число нейронов предыдущего слоя (или, грубо говоря, как раз ширина \(d\)).
И в стандартной параметризации, и в μP инициализация скрытых слоев масштабируется одинаково — как \(1 / \sqrt{\text{fan\_in}}\): больше входов, значит, каждый вклад меньше, значит, сумма остаётся того же порядка.
Для скрытых слоев μP меняет не инициализацию, а learning rate (см. Правило 2). А более агрессивное масштабирование инициализации (как 1/fan_in) относится к выходному слою, а не к скрытым.
2. Learning rate. LR скрытых слоёв масштабируется как \(1 / m_\text{width}\), где
\[ m_\text{width} = \frac{d_\text{модели}^{\text{целевой}}}{d_\text{модели}^{\text{базовой}}} \]То есть если прокси-модель имеет \(d = 256\), а целевая \(d = 4096\), то \(m_\text{width} = 16\), и LR скрытых слоев целевой модели будет в 16 раз меньше, чем у прокси.
Для Adam-оптимизатора LR скрытых слоев \(\propto 1/m_\text{width}\). Входной слой (embedding) и выходной слой (logits) трактуются отдельно, у них свои масштабы. Отдельно масштабируются и attention-логиты: в μP их делят на d_head, а не на \(\sqrt{\text{d\_head}}\), как в обычном attention (это учитывает корреляцию Q и K, которая возникает при обучении). Но это детали, а общий смысл в том, что мы фиксируем масштаб обновлений.
Итог: оптимальные гиперпараметры перестают зависеть от ширины. Их можно найти на маленькой модели и перенести на большую без повторного поиска.
Что переносится, а что нет
Чтобы не создавать ложного впечатления, что μP переносит вообще все, посмотрим, что может и не так чисто перенестись из маленькой модели в большую:
| переносится хорошо (по ширине) | не переносится / не гарантировано |
|---|---|
| Learning rate | Dropout, weight decay |
| Параметры инициализации | Длина обучения (число токенов) |
| Layer-wise множители LR | Все, что завязано на глубину |
| Архитектурные изменения |
Поэтому тут надо понимать ограничение: μP снимает зависимость оптимизационных гиперпараметров только от ширины.
Но при этом все, что живет по другим осям (глубина, объем данных, регуляризация), под его гарантию не попадает.
Dropout и weight decay — это про то, сколько и как модель учится, а не про масштаб обновлений, поэтому их все равно придется подбирать отдельно.
02Пример: как выглядит перенос
Подобрали на прокси-модели шириной \(d = 256\) оптимальный LR = 0.003 (как подбирали — см. ниже, раздел про sweep). Хотите обучить модель шириной \(d = 4096\).
Считаем: \(m_\text{width} = 4096 / 256 = 16\). Модель в 16 раз шире.
Под μP:
- LR скрытых слоёв целевой модели = \(0.003 / 16 \approx 1.9 \times 10^{-4}\)
- Разброс инициализации сжимается в \(\sqrt{16} = 4\) раза
И все. Никакого перебора на целевой модели. Никаких 200 прогонов на \(d = 4096\). Одна формула, один расчет, и мы стартуем с LR, который с высокой вероятностью достаточно близок к оптимальному.
Без μP нам пришлось бы гонять полный поиск LR прямо на 4096-модели. На таком масштабе один прогон — это дни GPU-времени и десятки тысяч долларов. И таких прогонов нужны были бы десятки.
03Как находят LR на маленькой модели
Нахождение LR на маленькой модели представляет из себя обычный перебор, который на английском называют sweep.
Берем маленькую модель (прокси).
Запускаем ее обучение много раз с разными learning rate: скажем, \(10^{-4}\), \(3 \times 10^{-4}\), \(10^{-3}\), \(3 \times 10^{-3}\), \(10^{-2}\)…
Каждый прогон достаточно дешев, потому что модель крошечная — миллионы параметров вместо миллиардов, и мы ее обучаем не на триллионе токенов, а на нескольких миллионах. Прогон занимает минуты или часы, не недели.
Далее смотрим, при каком LR итоговый loss минимальный (или кривая обучения стабильнее всего). Это и будет нашим оптимум для данной ширины.
Когда в статьях пишут «200 запусков» — это и есть sweep: 200 дешевых и быстрых прогонов маленькой модели с разными комбинациями гиперпараметров (LR, init scale, иногда layer-wise множители), чтобы найти лучшую комбинацию.
На большой модели такой перебор невозможен — 200 прогонов даже какой-нибудь старой GPT-3 разорят кого угодно.
Вся идея μP: сделать этот дорогой перебор один раз на дешевой модели, а потом перенести найденное на большую по правилу масштабирования.
Для того, чтобы перенос сработал, обе модели — и прокси, и целевая — должны быть записаны в μP-параметризации. Если прокси обучалась в SP, а целевая в μP, формулы масштабирования не применимы напрямую. Обе должны быть в одной системе координат.
04Результат: 6.7B бьёт 13B
В статье Yang et al. (2022) есть конкретный эксперимент, который лучше всего показывает, зачем всё это нужно. Три модели, и у каждой своя роль:
| модель | роль в эксперименте |
|---|---|
| 40M | Прокси. На ней гоняли перебор LR (те самые ~200 запусков). Нужна только чтобы найти хороший LR дешево. |
| 6.7B | Целевая модель. На нее μTransfer перенес найденный LR. Ее обучили один раз, с хорошим LR, без всякого перебора. |
| 13B | Точка отсчета. Существующая модель GPT-3 такого размера. Приводится, чтобы сказать: «вот такого уровня качества достигла наша 6.7B». |
Ключевой результат: 6.7B, обученная с перенесенным LR, дотянулась по качеству (loss на валидации) до 13B, обученной без μP. Модель вдвое меньшего размера показала качество модели вдвое большего размера.
Откуда берется двукратная экономия в этой конкретной статье
Тут важно не запутаться: экономия не в том, что подбор LR стал дешевле, а в итоговом размере модели, который нужен ради заданного качества.
Цепочка такая:
- Хороший LR (перенесенный с 40M) выжал из 6.7B заметно больше качества, чем она дала бы с посредственным LR. Настолько больше, что 6.7B догнала 13B.
- Значит, чтобы получить качество уровня 13B модели, хватило модели вдвое меньше — 6.7B вместо 13B.
- Модель вдвое меньше и обучать примерно вдвое дешевле (меньше FLOPs, меньше GPU-часов, меньше времени).
То есть тут история двойная:
- сначала мы дешево подобрали LR по методу, который разобрали выше;
- затем, как раз из-за того, что подобрали мы его неприлично хорошо, мы смогли обучить большую модель среднего размера до качества большой модели в два раза больше.
А это значит, что мы не только сэкономили в процессе обучения через экономию за подбор LR, но и то, что мы бесконечно будем экономить на инференсе более мелкой модели того же качества.
????? PROFIT!!!!!11
05Почему μP говорит именно про ширину
У трансформера два разных измерения размера:
- Глубина — это то количество слоев или блоков, которые поставлены друг за другом. Данные проходят через них последовательно, слой за слоем.
- Ширина — это та величина, который показывает, насколько толст каждый слой из слоев. В первую очередь это d_model — длина вектора, которым представлен каждый токен, плюс связанные размерности (скрытый слой в FFN, число голов внимания).
Если модель — это труба, по которой текут данные, то глубина — это длина трубы, а ширина — ее диаметр.
Почему параметры растут в основном от ширины
Хорошее приближение для числа параметров трансформера:
\[ N \approx 12 \cdot L \cdot d^2 \]Откуда берется \(12d^2\) на слой?
В attention четыре матрицы размера \(d \times d\) (Q, K, V, выход) = \(4d^2\), в FFN две матрицы \(d \times 4d\) = \(8d^2\).
Получается \(\sim 12d^2\) на слой, умножить на \(L\) слоев.
Ключевое здесь — \(d\) в квадрате:
- Углубить (\(\times\) слоев) → параметры растут линейно: вдвое слоев = вдвое параметров.
- Расширить (\(\times d_\text{model}\)) → параметры растут квадратично: вдвое шире = вчетверо параметров.
Поэтому, когда модель растет с 124M до 405B, основная часть роста параметров сделана как раз за счет роста вширь, а не вглубь.
Почему математика μP работает именно вдоль ширины
μP выведен в пределе бесконечной ширины (\(d_\text{model} \to \infty\)).
В этом пределе все масштабируется красиво и предсказуемо: рост числа слагаемых в суммах, рост размерности матриц — все компенсируется аккуратным подбором масштабов.
С глубиной все гораздо сложнее.
Там нет такого чистого предела, и перенос гиперпараметров работает хуже.
Были попытки сделать μP по глубине, но это уже не тот аккуратный результат, который мы рассматриваем сегодня, а скорее некий набор получившихся эмпирических правил.
Так что μP в базовом виде — это только перенос между узкой и широкой версией модели при той же глубине.
Реальные конфиги: ширина и глубина в цифрах
Лестница от 124M до 405B:
| модель | количество параметров | количество слоев (L) | ширина (d_model) | количество голов | d_model / L |
|---|---|---|---|---|---|
| GPT-2 Small | 124M | 12 | 768 | 12 | 64 |
| GPT-3 Medium | 350M | 24 | 1024 | 16 | 43 |
| GPT-3 XL | 1.3B | 24 | 2048 | 24 | 85 |
| GPT-3 6.7B | 6.7B | 32 | 4096 | 32 | 128 |
| GPT-3 13B | 13B | 40 | 5120 | 40 | 128 |
| Llama-3 70B | 70B | 80 | 8192 | 64 | 102 |
| GPT-3 175B | 175B | 96 | 12288 | 96 | 128 |
| Llama-3 405B | 405B | 126 | 16384 | 128 | 130 |
У больших моделей соотношение ширины к длине \(d_\text{model} / L\) уже долгое время остается примерно постоянным, и почти весь бюджет роста параметров инвестируется именно в рост ширины.
Раньше модели делали тонкими и длинными. Вообще, многие года примерно с 2015 по 2020 мы по большей части боролись за глубину. С ней ассоциировалась мощь и качество модели.
В последние же года при одинаковом объеме параметров, новые модели стали значительно мельче и шире.
Например обратите внимания на переход:
- GPT-2 XL (1.5B) состояла из 48 слоёв × 1600 ширины;
- GPT-3 XL (1.3B) состояла уже из 24 слоя × 2048 ширины, вдвое мельче по глубине и шире.
Тот же класс размера, а при проектировании GPT-3 ребята сознательно выбрали меньшее количество слоев и большую ширину.
При этом сама глубина с ростом масштаба все равно увеличивается. В таблице L растет с 12 до 126, просто ширина при этом все равно растет быстрее.
06Ограничения и нюансы
- Аккуратность реализации
В процессе применения техники очень легко ошибиться в том, какой слой каким масштабом масштабировать. Если где-то получилась ошибка, то весь перенос ломается, и мы получаем мусорный LR на целевой модели. В открытых реализациях, например, mup-репозитории от Microsoft, — это уже реализовано, но если делать свою реализацию с нуля руками, нужно быть очень аккуратным. - Перенос дает стартовую точку, не абсолют
Иногда LR все равно нужно чуть-чуть докрутить уже на целевом масштабе, но это все равно всего 2–3 прогона вместо 200. Экономия на порядки, но не совсем ноль работы. - Прокси и цель должны совпадать по всему, кроме ширины
Тот же датасет, та же глубина, тот же оптимизатор, та же последовательность токенов. Чем сильнее прокси отличается от цели чем-то еще, кроме ширины, тем менее точным получается перенос. Если вы меняется одновременно ширина, число слоев и датасет, то никакой μP не поможет. - Не работает по глубине
Если целевая модель не только шире, но и глубже, перенос по ширине все еще применим, но глубинную часть придется подбирать отдельно или полагаться на эвристики.