Языковая модель Pathways (PaLM) — это плотно активированная языковая модель Transformer с 540 миллиардами параметров. Он обучается на чипах 6144 TPU v4 с использованием Pathways, новой системы машинного обучения, которая обеспечивает высокоэффективное обучение на нескольких модулях TPU.

PaLM 540B достигает прорывной производительности, превосходя передовые технологии в наборе многошаговых логических задач и превосходя среднюю производительность человека в недавно выпущенном тесте BIG-bench. Значительное количество задач BIG-bench показало прерывистые улучшения по сравнению с масштабом модели, а это означает, что производительность резко возросла по мере масштабирования до самой большой модели. PaLM также обладает широкими возможностями в многоязычных задачах и генерации исходного кода.

Архитектура модели

PaLM использует стандартную архитектуру модели Transformer в конфигурации только с декодером со следующими изменениями:

Активация SwiGLU

Активации SwiGLU (Swish(xW) · xV) используются для промежуточных активаций MLP, поскольку было показано, что они значительно повышают качество по сравнению со стандартными активациями ReLU, GeLU или Swish.

Параллельные слои

В каждом блоке Transformer используется «параллельная» формулировка, а не стандартная «последовательная» формулировка. В частности, стандартная формулировка может быть записана как:

y = x + MLP(LayerNorm(x + Attention(LayerNorm(x)))

Принимая во внимание, что параллельная формулировка может быть записана как:

y = x + MLP (Норма Слоя (x)) + Внимание (Норма Слоя (x))

Параллельная формулировка приводит к увеличению скорости обучения примерно на 15% в больших масштабах, поскольку умножения MLP и входной матрицы внимания могут быть объединены. Эксперименты по абляции показали небольшое ухудшение качества по шкале 8B, но не ухудшение качества по шкале 62B, что привело к экстраполяции того, что эффект параллельных слоев должен быть нейтральным по качеству при шкале 540B.

Внимание с несколькими запросами

Стандартная формулировка Transformer использует k головок внимания, где входной вектор для каждого временного шага линейно проецируется на тензоры «запрос», «ключ» и «значение» формы [k, h]. , где h — размер головы, привлекающей внимание. Здесь проекции ключ/значение являются общими для каждой головы, т. е. «ключ» и «значение» проецируются на [1, h], но «запрос» по-прежнему проецируется на форму [ k, h]. Это нейтрально влияет на качество модели и скорость обучения, но приводит к значительной экономии средств во время авторегрессионного декодирования. Это связано с тем, что стандартное многоголовое внимание имеет низкую эффективность на аппаратном ускорителе во время авторегрессивного декодирования, тензоры ключ/значение не разделяются между примерами, и за раз декодируется только один токен.

Веревочные заделки

Встраивания RoPE используются, а не вложения абсолютного или относительного положения, поскольку было показано, что вложения RoPE имеют лучшую производительность при длинных последовательностях.

Общие вложения ввода-вывода

Матрицы встраивания ввода и вывода являются общими, что часто (но не всегда) делалось в прошлой работе.

Нет предубеждений

Никакие смещения не использовались ни в одном из плотных ядер или норм слоя. Было обнаружено, что это привело к повышению стабильности обучения для больших моделей.

Словарный запас

Используется словарь SentencePiece с 256 тыс. токенов, который был выбран для поддержки большого количества языков в обучающем корпусе без чрезмерной токенизации. Словарь полностью без потерь и обратим, что означает, что пробелы полностью сохраняются в словаре (особенно важно для кода), а символы Unicode вне словаря разбиваются на байты UTF-8 с токеном словаря для каждого байта. Числа всегда разбиваются на отдельные токены цифр.

Гиперпараметры масштаба модели

Набор обучающих данных

Набор данных для предварительной подготовки PaLM состоит из высококачественного корпуса из 780 миллиардов токенов, представляющих широкий спектр вариантов использования естественного языка. Набор данных представляет собой смесь отфильтрованных веб-страниц, книг, статей в Википедии, новостных статей, исходного кода и разговоров в социальных сетях. Этот набор данных основан на наборах данных, используемых для обучения LaMDA и GLaM.

Все три модели обучаются ровно на одной эпохе данных (перемешиваются одинаково для всех моделей), а пропорции смешивания выбираются таким образом, чтобы избежать повторения данных в любом подкомпоненте.

В дополнение к данным на естественном языке в набор данных для предварительной подготовки включен код. Исходный код в наборе данных для предварительной подготовки получен из репозиториев с открытым исходным кодом на GitHub. Файлы фильтруются по лицензии, включенной в репозиторий, за исключением лицензий с авторским левом. Файлы также фильтруются по расширению имени файла, чтобы ограничить их одним из 24 распространенных языков программирования, включая Java, HTML, JavaScript, Python, PHP, C#, XML, C++ и C, что приводит к 196 ГБ исходного кода. Кроме того, дубликаты удаляются на основе расстояния Левенштейна между файлами, поскольку известно, что дубликаты файлов распространены в репозиториях исходного кода.

Учебная инфраструктура

Модель обучается с использованием фреймворков JAX и T5X на модулях TPU v4, которые состоят из чипов TPU, подключенных через сеть центра обработки данных. Система, используемая в этом случае, включает 3072 чипа TPU v4, подключенных к 768 хостам, что обеспечивает эффективное масштабирование без параллелизма конвейера.

В предыдущих отчетах о крупномасштабном обучении использовались разные подходы. Некоторые модели обучались на одной системе TPU без конвейерного параллелизма или использования сети центра обработки данных, в то время как в других использовалась комбинация модели, данных и конвейерного параллелизма или конвейерная обработка между модулями.

Существуют определенные недостатки конвейерной обработки, такие как время простоя и более высокие требования к пропускной способности памяти. Чтобы преодолеть эти ограничения, PaLM использовала стратегию без конвейерного параллелизма. Каждый модуль TPU v4 содержит полную копию параметров модели, разделенную на несколько микросхем с использованием параллелизма модели и параллелизма данных с полным сегментированием. Во время прямого прохода веса собираются по параллельной оси данных, и сохраняются выбранные тензоры активации. Во время обратного прохода оставшиеся активации повторно материализуются, что приводит к повышению производительности обучения при больших размерах пакетов.

Для масштабирования обучения за пределы одного модуля TPU v4 используется система Pathways. Он использует архитектуру клиент-сервер для достижения двустороннего параллелизма данных на уровне модуля. Клиент Python отправляет половину обучающего пакета в каждый модуль, где прямые и обратные вычисления выполняются параллельно с использованием данных внутри модуля и параллелизма моделей. Градиенты, вычисленные в каждом модуле, затем передаются в другой модуль, а обновления параметров применяются параллельно для получения идентичных параметров для следующего временного шага.

Клиент Python создает программу сегментированного потока данных, которая запускает работу JAX/XLA на удаленных серверах, содержащих модули TPU. Программа включает компоненты для вычислений внутри модуля, передачи градиента между модулями и обновления оптимизатора. Задержка маскируется с помощью асинхронного группового планирования, а стоимость управления передачей данных амортизируется с помощью модели выполнения потока данных с сегментированием.

Существуют определенные проблемы с достижением высокой пропускной способности обучения для передачи градиента между модулями в масштабе 6144 чипов TPU v4. Чтобы решить эту проблему, сетевой стек Pathways тщательно разработан, разбивая градиентные передачи на более мелкие фрагменты и направляя их через несколько меньших потоков по разным сетевым каналам центра обработки данных. С этими оптимизациями пропускная способность обучения примерно в 1,95 раза выше по сравнению с одним модулем во время обучения с некоторыми пробелами в производительности из-за отсутствия перекрытия между обратным проходом и уменьшением градиента между модулями.

Эффективность обучения

Использование аппаратных FLOP (HFU) зависит от системы, может варьироваться в зависимости от реализации и выбора конструкции и не учитывает конечной цели достижения высокой пропускной способности в токенах в секунду. Следовательно, используется новая метрика, называемая моделью FLOPs Utilization (MFU), которая не зависит от реализации и позволяет более справедливо сравнивать эффективность между различными системами.

Настройка обучения

Инициализация веса

Веса ядра (т. е. все, кроме вложений и шкал норм слоев) инициализируются с помощью «масштабирования дисперсии веером», т. е. W ∼ N (0, 1/√ nin), где nin — входная размерность ядра. Входные вложения инициализируются как E ∼ N (0, 1), поскольку к вложениям не применяется нормализация слоя. Поскольку входной и выходной слои внедрения являются общими, мы масштабируем выходные логиты до softmax на 1/√n, где n — размер внедрения.

Оптимизатор

Модель обучалась с помощью оптимизатора Adafactor, без факторизации. Это эффективно эквивалентно Адаму с «масштабированием параметров», которое масштабирует скорость обучения на среднеквадратичное значение матрицы параметров. Поскольку инициализация веса пропорциональна 1/√n, эффект от этого аналогичен ручному уменьшению скорости обучения Адама.

Гиперпараметры оптимизации

Скорость обучения Adafactor, равная 10–2, используется для первых 10 000 шагов, которые затем затухают со скоростью 1/√k, где k — номер шага. Для обучения используются β1 = 0,9 и β2 = 1,0 − k − 0,8. Было обнаружено, что это более стабильно, чем стандартное β2 = 0,99, при обучении больших языковых моделей, потому что плохая оценка секундных моментов в более коротких окнах может быть вызвана редкими встраивающимися токенами. Отсечение градиента глобальной нормы со значением 1,0 используется для всех моделей. Во время обучения используется динамическое уменьшение веса lr², где lr — текущая скорость обучения.

Функция потерь

Модель обучается с помощью функции потерь моделирования стандартного языка, которая представляет собой среднюю логарифмическую вероятность всех токенов без сглаживания меток. Вспомогательная потеря z loss = 10 ^ −4 · log2 Z дополнительно используется для того, чтобы нормализатор softmax log(Z) был близок к 0, поскольку было обнаружено, что он повышает стабильность обучения.

Длина последовательности

Длина последовательности 2048 использовалась для всех моделей. Входные примеры объединяются вместе, а затем разбиваются на последовательности ровно из 2048 токенов, так что нет дополнительных токенов, но примеры могут быть разделены посередине. Входные примеры отличаются друг от друга специальным токеном [eod].

Размер партии

Для всех моделей размер пакета увеличивается во время обучения. Для самой большой модели размер пакета 512 (токены 1M) используется до шага 50k, затем он удваивается до 1024 (токенов 2M) до шага 115k и, наконец, снова удваивается до 2048 (токенов 4M) до завершения обучения на шаге 255k. Аналогичные графики использовались для меньших моделей. Причина использования такого графика размера партии двояка: (1) меньшие размеры партии более эффективны для выборки (т. е. лучше потери в зависимости от наблюдаемых токенов) на ранних этапах обучения, в то время как большие размеры партий выгодны на более поздних этапах обучения из-за лучшего оценки градиента, и (2) большие размеры пакетов приводят к большим размерам матричного умножения, что повышает эффективность TPU.

Побитовый детерминизм

Модель полностью побитово воспроизводима с любой контрольной точки. Другими словами, если модель была обучена до 17 000 шагов за один прогон и мы перезапустили ее с контрольной точки 15 000, то система обучения гарантированно даст идентичные результаты в обоих прогонах с контрольной точки 15 000 до 17 000. Это достигается двумя способами: (1) побитово-детерминистической структурой моделирования, предоставляемой JAX+XLA+T5X, и (2) конвейером детерминированного набора данных, где перетасованные данные записываются в формате с произвольным доступом, поэтому содержимое данная тренировочная партия является только функцией номера шага.

Выбывать

Модель обучалась без отсева, хотя в большинстве случаев для точной настройки используется отсев 0,1.

Тренировочная нестабильность

Во время обучения самой большой модели часто возникали всплески функции потерь, несмотря на то, что было реализовано отсечение градиента. Эти всплески происходили через неравные промежутки времени, иногда в конце тренировочного процесса, и не наблюдались в моделях меньшего размера. Из-за высокой стоимости обучения самой большой модели не удалось определить четко определенный подход к устранению этих всплесков.

Однако была найдена простая стратегия, позволяющая эффективно смягчить эту проблему. Обучение было перезапущено с контрольной точки примерно за 100 шагов до того, как произошел всплеск, и было пропущено около 200–500 пакетов данных, включая пакеты, наблюдаемые до и во время всплеска. При использовании этого метода смягчения потери не увеличились снова в тот же момент.

Считалось, что всплески не вызваны ошибочными данными, поскольку эксперименты проводились с использованием окружающих пакетов данных, обучением с другой контрольной точки, и в этих случаях не наблюдалось всплесков. Это говорит о том, что всплески возникают только из-за определенных комбинаций пакетов данных и состояния параметров модели.

Оценка

Задачи НЛП на английском языке

PaLM 540B превосходит предыдущие модели SOTA, превосходя их в 24 из 29 задач с одним выстрелом и в 28 из 29 задач с несколькими выстрелами. Примечательно, что он достигает значительного улучшения на 10+ баллов в задачах на понимание прочитанного и NLI. Не ограничиваясь размером, PaLM 540B даже превосходит сопоставимую модель (Megatron-Turing NLG 530B) по всем тестам, подчеркивая важность данных предварительной подготовки, стратегии обучения и количества токенов.

Массивное многозадачное понимание языка

PaLM 540B улучшает средний балл теста MMLU примерно на 2 балла. и превосходит шиншиллу во всех категориях, кроме одной.

Тонкая настройка

Эксперименты по тонкой настройке модели PaLM проводились на тесте SuperGLUE. PaLM был настроен с помощью оптимизатора Adafactor со скоростью обучения 5 × 10–5 и размером пакета 32. Как правило, сходимость достигается с помощью PaLM менее чем за 15 000 шагов тонкой настройки.

PaLM обеспечивает конкурентоспособную производительность, близкую к SOTA, даже по сравнению с моделями кодер-декодер, которые, как показано, работают лучше, чем модели только с декодером.

Существует значительный разрыв между результатами нескольких выстрелов и результатами точной настройки.

PaLM конкурирует с самыми современными моделями (модели кодер-декодер) и значительно превосходит модели только с декодером.

БОЛЬШАЯ скамья

PaLM значительно превосходит GPT-3, Gopher и Chinchilla, а PaLM 540B с 5 выстрелами набирает более высокий балл, чем средний балл людей, которых просят решить те же задачи.

PaLM 540B 5-shot превосходит предыдущую SOTA в 44 из 58 стандартных задач.

PaLM 540B превосходит среднюю человеческую производительность в совокупности, она все еще выше, чем PaLM 540B в 35% отдельных задач.

Рассуждение

Используя 8-шаговую цепочку мыслей в сочетании с внешним калькулятором, PaLM 540B достигает производительности 58%, превосходя предыдущую SOTA с 55% от GPT-3.

В 7 наборах данных рассуждений 8-шаговый прогноз с PaLM 540B+цепочка размышлений достиг точности SOTA в 4 задачах и близок к SOTA в оставшихся 3 задачах.

Задачи кода

Модель PaLM обеспечивает сравнимую производительность в нескольких оценках с ранее опубликованными результатами, используя в 50 раз меньше кода Python. Это может быть связано с комбинацией переноса из других языков программирования и данных на естественном языке, а также с тем, что более крупные модели могут быть более эффективными с точки зрения выборки, чем более мелкие модели.

Многоязычная генерация естественного языка

PaLM 540B с тонкой настройкой соответствует или улучшает лучшую SOTA во всех задачах генерации на английском языке, достигая нового современного результата тонкой настройки для 4 из 6 задач суммирования не на английском языке.

Многоязычный ответ на вопрос

PaLM 540B достигает конкурентоспособных результатов, несмотря на то, что не тренируется на таком большом количестве неанглоязычных данных.

Бумага

PaLM: масштабируемое языковое моделирование с путями 2204.02311

Хотите узнать больше?

Не пропустите другие увлекательные темы из этой серии. Просто нажмите здесь и откройте для себя новейшее исследование!

Подпишитесь на еженедельные обновления!!