Training Compute-Optimal Large Language Models

Jordan Hoffmann; Sebastian Borgeaud; Arthur Mensch; Elena Buchatskaya; Trevor Cai; Eliza Rutherford; Diego de Las Casas; Lisa Anne Hendricks; Johannes Welbl; Aidan Clark; Tom Hennigan; Eric Noland; Katie Millican; George van den Driessche; Bogdan Damoc; Aurelia Guy; Simon Osindero; Karen Simonyan; Erich Elsen; Jack W. Rae; Oriol Vinyals; Laurent Sifre

Статьи · arXiv 2203.15556

Обучение вычислительно-оптимальных больших языковых моделей

Training Compute-Optimal Large Language Models

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland et al.

Опубликовано 2022-03-29cs.CLcs.LG arXiv PDF

Читать с переводом и ИИ

Кратко в трёх строках

🎯 Определение оптимального баланса между параметрами модели и обучающими токенами при фиксированном бюджете вычислений.
🔧 Эмпирическая аппроксимация законов масштабирования на основе >400 прогонов моделей Transformer с помощью огибающих минимального лосса, срезов IsoFLOP и параметрического моделирования функции потерь.
📈 Параметры и токены должны масштабироваться одинаково ( $a \approx 0.5$ , $b \approx 0.5$ ). Модель Chinchilla (70B), обученная на 1.4T токенов, укладывается в бюджет вычислений Gopher и достигает 67.6% на MMLU.

Аннотация

Мы исследуем оптимальный размер модели и количество токенов для обучения языковой модели на архитектуре Transformer при заданном вычислительном бюджете. Мы показываем, что современные большие языковые модели существенно недообучены — это результат недавней тенденции масштабировать модели, сохраняя фиксированным объем обучающих данных. Обучив более 400 языковых моделей размером от 70 миллионов до более чем 16 миллиардов параметров на объемах от 5 до 500 миллиардов токенов, мы выяснили, что для вычислительно-оптимального обучения размер модели и количество токенов должны масштабироваться одинаково: с каждым удвоением размера модели количество обучающих токенов также должно удваиваться. Мы проверяем эту гипотезу, обучив предсказанную вычислительно-оптимальную модель Chinchilla, которая использует тот же вычислительный бюджет, что и Gopher, но имеет 70B параметров и в 4 раза больше данных. Chinchilla стабильно и существенно превосходит Gopher (280B), GPT-3 (175B), Jurassic-1 (178B) и Megatron-Turing NLG (530B) на широком спектре задач. Это также означает, что Chinchilla требует существенно меньше вычислений для файнтюнинга и инференса, значительно упрощая практическое применение. В частности, Chinchilla достигает рекордной средней точности 67.5% на бенчмарке MMLU, что более чем на 7% превосходит результат Gopher.

Оригинальная аннотация (англ.)

We investigate the optimal model size and number of tokens for training a transformer language model under a given compute budget. We find that current large language models are significantly undertrained, a consequence of the recent focus on scaling language models whilst keeping the amount of training data constant. By training over 400 language models ranging from 70 million to over 16 billion parameters on 5 to 500 billion tokens, we find that for compute-optimal training, the model size and the number of training tokens should be scaled equally: for every doubling of model size the number of training tokens should also be doubled. We test this hypothesis by training a predicted compute-optimal model, Chinchilla, that uses the same compute budget as Gopher but with 70B parameters and 4 $\times$ more more data. Chinchilla uniformly and significantly outperforms Gopher (280B), GPT-3 (175B), Jurassic-1 (178B), and Megatron-Turing NLG (530B) on a large range of downstream evaluation tasks. This also means that Chinchilla uses substantially less compute for fine-tuning and inference, greatly facilitating downstream usage. As a highlight, Chinchilla reaches a state-of-the-art average accuracy of 67.5% on the MMLU benchmark, greater than a 7% improvement over Gopher.

Разбор статьи

Проблема и мотивация

Предыдущие работы (в частности, Kaplan et al., 2020) предполагали, что при увеличении вычислений следует агрессивно масштабировать параметры модели ( $N \propto C^{0.73}$ ), тогда как число токенов должно расти значительно медленнее ( $D \propto C^{0.27}$ ). Как следствие, такие крупные модели, как GPT-3 (175B), Jurassic-1 (178B), Gopher (280B) и MT-NLG (530B), практически всегда обучались приблизительно на 300 миллиардах токенов. Однако это привело к их существенному недообучению относительно затраченных вычислительных ресурсов. В данной работе этот баланс пересматривается: как при фиксированном бюджете вычислений $C$ распределить число параметров $N$ и число токенов $D$ , чтобы минимизировать потери на валидации?

Ключевые идеи

Равное масштабирование: при каждом удвоении параметров модели количество обучающих токенов также должно удваиваться ( $N_{o pt} \propto C^{0.5}$ , $D_{o pt} \propto C^{0.5}$ ).
Согласование косинусного расписания: расписание шага обучения (learning rate) должно точно соответствовать реальной длительности обучения. Использование промежуточных точек длинного запуска (например, с затуханием на 130B токенов) завышает потери для меньших объемов токенов и искажает эмпирические оценки масштабирования.
Эффективность в прикладных задачах: обучение меньшей модели на значительно большем числе токенов дает ту же или более высокую точность при равных вычислительных затратах на предобучение, одновременно кардинально снижая требования к памяти, затраты на файнтюнинг и задержку при инференсе.

Метод (с ключевыми уравнениями)

Авторы оценивают более 400 моделей (от 70M до 16B параметров, обученных на выборках от 5B до 500B токенов), применяя три взаимодополняющих подхода:

Минимум по кривым обучения: варьируется длительность обучения (4 варианта) для фиксированных архитектур. Выделяется огибающая минимального лосса на единицу FLOP и аппроксимируются степенные законы: $N_{o pt} \propto C^{a}, D_{o pt} \propto C^{b}$ что дает $a = 0.50$ , $b = 0.50$ .
Профили IsoFLOP: фиксируется общий объем FLOP для 9 дискретных бюджетов (от $6 \times 1 0^{18}$ до $3 \times 1 0^{21}$ ) с косинусным затуханием, согласованным с финальной точкой. Точки минимума лосса аппроксимируются параболами по числу параметров, что дает $a = 0.49$ , $b = 0.51$ .
Параметрическая аппроксимация функции потерь: сглаженный лосс моделируется через декомпозицию риска: $\hat{L} (N, D) ≜ E + \frac{A}{N ^{α}} + \frac{B}{D ^{β}}$ где $E$ обозначает энтропию естественного текста, $A / N^{α}$ отражает субоптимальность емкости модели, а $B / D^{β}$ — неоптимальность ранней остановки и ограниченности данных. Минимизация потерь Хьюбера на логарифмах невязок дает $E = 1.69$ , $A = 406.4$ , $B = 410.7$ , $α = 0.34$ , $β = 0.28$ . При ограничении $FLOPs (N, D) \approx 6 N D$ : $a = \frac{β}{α + β} = 0.46, b = \frac{α}{α + β} = 0.54$

Результаты (с ключевыми цифрами)

Chinchilla 70B: обучена с тем же бюджетом FLOP, что и Gopher ( $5.76 \times 1 0^{23}$ FLOPs, или $6.3 \times 1 0^{23}$ при точном послойном учете), но имеет 70B параметров и 1.4 триллиона токенов.
Бенчмарк MMLU: Chinchilla достигает средней точности 67.6% в режиме 5-shot, опережая Gopher (60.0%), GPT-3 (43.9%) и превосходя прогноз экспертов-людей на июнь 2023 года (63.4%). Модель достигает точности $> 90%$ на 4 отдельных задачах MMLU.
Прикладные бенчмарки: превосходит Gopher на BIG-bench (65.1% против 54.4%), LAMBADA (77.4% против 74.5%) и closed-book Natural Questions (31.5% 5-shot против 24.5%).
Смещения и токсичность: спонтанная токсичность (unprompted toxicity) остается на одном уровне у Gopher и Chinchilla (медиана 0.064 против 0.066), при этом общая точность на Winogender увеличивается с 71.4% до 78.3%.

Ограничения и открытые вопросы

Из-за высокой стоимости обучения напрямую были сопоставлены только две крупномасштабные передовые модели при равных вычислительных ресурсах (Gopher 280B против Chinchilla 70B).
Анализ исходит из обучения в одну эпоху на веб-корпусах; поведение при повторном использовании данных в течение нескольких эпох не изучалось.
Эмпирически наблюдается небольшая отрицательная кривизна (вогнутость) функции $lo g (N_{o pt})$ при больших вычислительных бюджетах: это говорит о том, что степенные законы могут все еще слегка переоценивать оптимальный размер модели.

Ключевые термины

Chinchilla: Авторегрессионная языковая модель на 70 миллиардов параметров, обученная на 1.4 триллиона токенов в рамках того же вычислительного бюджета, что и Gopher.
Gopher: Языковая модель Transformer на 280 миллиардов параметров, обученная на 300 миллиардах токенов, выступающая в качестве вычислительно-эквивалентного бейзлайна.
IsoFLOP profiles: Методология, при которой число параметров модели варьируется при постоянном общем количестве FLOPs для нахождения размера параметров с минимальными потерями.
Compute-optimal frontier: Математическая кривая, описывающая зависимость между доступным объемом FLOPs и оптимальным распределением числа параметров и количества обучающих токенов.
Parametric loss function: Статистическая декомпозиция потерь языковой модели на неустранимую энтропию Байеса, дефицит емкости параметров и ошибку оптимизации из-за конечного объема токенов.
Huber loss: Робастная функция потерь, использованная в подходе 3 для аппроксимации эмпирических запусков языковых моделей с ослаблением влияния выбросов на малых масштабах вычислений.
MassiveText: Многодоменный корпус предобучения, содержащий веб-страницы, книги, новости, программный код и Википедию, использованный для обучения как Gopher, так и Chinchilla.
FLOPs calculation: Оценка объема операций умножения с накоплением, приблизительно рассчитываемая как $6 N D$ для прямого и обратного проходов или через послойный учет операций.

Вопросы

Почему Kaplan et al. пришли к выводу, что параметры должны расти быстрее токенов?

Kaplan et al. использовали фиксированное косинусное расписание learning rate (например, рассчитанное на 130B токенов) и брали промежуточные несошедшиеся чекпоинты для меньших объемов токенов. Поскольку косинусное расписание теряет эффективность, если длина цикла существенно превышает фактическое число шагов, эти промежуточные потери оказались завышенными, что исказило оценку эффективности обучения меньших моделей на меньшем числе токенов.

Каково эмпирическое правило для вычислительно-оптимального масштабирования согласно Chinchilla?

Число параметров модели и количество обучающих токенов следует масштабировать в равных пропорциях: при увеличении объема вычислений в $2 \times$ раза размер модели должен расти в $\approx 2^{0.5} \approx 1.41 \times$ раза, а количество обучающих токенов — также в $\approx 2^{0.5} \approx 1.41 \times$ раза.

Сколько токенов рекомендуется для обучения вычислительно-оптимальной модели на 175B параметров?

Согласно оценкам статьи, модели на 175 миллиардов параметров требуется от 3.7 до 4.3 триллионов токенов и от $3.85 \times 1 0^{24}$ до $4.54 \times 1 0^{24}$ FLOPs, что многократно превышает 300 миллиардов токенов, использованных при обучении GPT-3.

Какие изменения в архитектуре были сделаны при переходе от Gopher к Chinchilla?

Помимо сокращения числа слоев и скрытой размерности с 280B до 70B, в Chinchilla оптимизатор Adam был заменен на AdamW, мастер-веса в распределенном состоянии оптимизатора сохранялись в float32, а также была отключена нормализация NFKC в токенизаторе SentencePiece.

Обзор на этой странице написан ИИ по тексту статьи и может содержать ошибки. Подробности смотрите в самой статье.

Читать с переводом и ИИ