Pythia: набор моделей для анализа больших языковых моделей в процессе обучения и масштабирования

Pythia - это набор языковых моделей, созданных скорее для изучения, чем для практического применения. Стелла Бидерман (Stella Biderman), Хейли Шёлькопф (Hailey Schoelkopf) и их коллеги из EleutherAI опубликовали статью в апреле 2023 года (arXiv 2304.01373, ICML 2023), а вычислительные ресурсы предоставила Stability AI. Набор включает 16 decoder-only моделей 8 размеров - от 70M до 12B параметров. Модели каждого размера обучались дважды: один раз на датасете Pile и один раз на его дедуплицированной версии, причём в каждом прогоне данные подавались в строго одинаковом порядке. В open-source-ai-reading-list Нейтана Ламберта этот отчёт стоит первым среди "полностью открытых", опережая линейку olmo от AI2, которая довела ту же идею до конкурентоспособного масштаба.

Зачем нужен набор и почему именно такой

Исходная посылка статьи: законы масштабирования описывают поведение уже обученных моделей по мере их роста, но почти ничто не связывает эти законы с тем, что происходит во время обучения, а немногочисленные попытки исследований опирались на закрытые модели или закрытые checkpoint'ы. Авторы сопоставили открытые семейства моделей начала 2023 года (GPT-2, GPT-3, GPT-Neo, OPT, T5, BLOOM) с тремя требованиями и выяснили, что ни одно из них не удовлетворяет всем трём: охват нескольких порядков масштаба, обучение на одних и тех же данных в одинаковом порядке, а также открытый доступ к данным и промежуточным checkpoint'ам.

В обзоре разбирается, чего не хватало конкурентам. GPT-3 получила половинную оценку за открытость, так как её API платный, а OpenAI ограничивает допустимые темы исследований. Модели GPT-Neo используют так, будто это единый набор, однако они различаются токенизатором, кодовой базой и порядком данных. Обучающие данные OPT закрыты, а в открытый доступ выложили менее 10 checkpoint'ов и только для трёх размеров. В BLOOM расходимость обучения устраняли откатом назад и пропуском проблемных последовательностей, из-за чего модели различаются мелкими недокументированными деталями. Если считать checkpoint'ы по модели с наименьшим их числом в наборе, лучшим конкурентом был GPT-Neo с 30 checkpoint'ами; Pythia же предлагает по 154 для каждой модели.

В Pythia намеренно пожертвовали производительностью ради контроля. Архитектура использует параллельные слои внимания и feedforward на всех размерах, хотя считалось, что такое решение вредит моделям меньше 6B. Авторы ожидали падения качества, но его не произошло: модели не уступают OPT при том же числе параметров и токенов.

Что именно обучалось

Все модели обучались на 299 892 736 000 токенах (около 300B), чтобы сравняться по числу токенов с исходными наборами GPT-3 и OPT. Стандартный Pile содержит 334B токенов по токенизатору GPT-NeoX, поэтому стандартные прогоны часть данных так и не увидели. Дедуплицированная копия (нечёткая дедупликация MinHashLSH с порогом 0.87) насчитывает всего 207B токенов, то есть дедуплицированные модели обучались около 1.5 эпох. В статье это превратили в экспериментальную возможность: сравнение checkpoint'ов непосредственно перед границей эпохи и сразу после неё не выявило никаких признаков того, что вторая эпоха ухудшает результаты на бенчмарках.

Архитектура повторяет GPT-3 с изменениями, ставшими стандартом к 2023 году: полносвязное внимание (dense attention) вместо чередующихся разреженных слоёв GPT-3, Flash Attention, ротационные эмбеддинги (RoPE, применённые к 25% измерений согласно таблице конфигурации), а также параллельная компоновка attention/MLP и инициализация в стиле GPT-J. Матрицы входных и выходных эмбеддингов разделены (untied), поскольку более ранние работы указывали, что это упрощает исследования интерпретируемости. Токенизатор - BPE, обученный на Pile для GPT-NeoX-20B.

Решение по обучению, сильнее всего идущее вразрез с общепринятыми представлениями, - это размер батча. Во всех моделях, включая вариант на 70M, используется 1024 последовательности по 2048 токенов - около 2M токенов на шаг. Считалось, что малым моделям для сходимости нужны маленькие батчи; авторы же не зафиксировали проблем со сходимостью при размере батча в 4-8 раз больше обычного для моделей до 1B, получив ускорение по астрономическому времени (wall-clock time) до 10 раз на малых размерах. Checkpoint сохраняется каждую 1000 шагов (143 штуки за 143 000 шагов), плюс шаг 0 и шаги с логарифмическим шагом 1, 2, 4 вплоть до 512, что даёт 154 checkpoint'а на модель. Обучение проходило на ускорителях A100 с 40GB памяти с использованием GPT-NeoX, DeepSpeed и ZeRO; общий объём вычислений для статьи, включая описанное ниже повторное обучение, составил 544 280 A100-часов.

В статье открыто говорится, что набор обучали дважды. В предварительном релизе "v0" размеры батчей были несогласованы (4M токенов для трёх размеров), а график скорости обучения (learning-rate schedule) спадал до 10% от пика для моделей до 2.8B, но до нуля - для 6.9B и 12B. Авторы переобучили всё с единообразными настройками, заново провели весь анализ и заменили публичные модели 31 марта 2023 года, сохранив v0 доступной для воспроизводимости. В январе 2023 года они также переименовали модели, чтобы число параметров включало слои эмбеддингов, как в GPT-2, BLOOM и OPT. Оценив затраты на альтернативный вариант с обучением множества случайных seed'ов для каждого размера, от него отказались: 25 seed'ов обошлись бы примерно в 10 миллионов долларов США.

Что было опубликовано

Всё выпущено под лицензией Apache 2.0: все веса и checkpoint'ы на HuggingFace Hub, код обучения, код оценки с сырыми результатами бенчмарков и данные. Публикация данных не ограничивается ссылкой на Pile. EleutherAI предоставляет предварительно токенизированные файлы, которые непосредственно читал dataloader, и скрипт, воссоздающий точный dataloader, чтобы исследователь мог извлечь содержимое любого батча на любом шаге. Именно на этом свойстве строятся практические исследования статьи, и именно оно отличает Pythia от обычного набора открытых весов.

Оценка проводилась с помощью собственного фреймворка EleutherAI LM Evaluation Harness, причём все метрики рассчитывались самостоятельно, а не копировались из других статей. На восьми распространённых бенчмарках Pythia и Pythia-deduped показывают результаты, очень близкие к OPT и BLOOM. Таблицы с результатами по отдельным бенчмаркам в приложении не сохранились при извлечении в локальную копию источника, как и таблица 1 с гиперпараметрами для каждого размера, поэтому точные оценки и количество слоёв на этой странице не приведены.

Три вывода, разошедшихся с литературой

Авторы выделяют три результата, противоречащих публикациям того времени. Дедупликация Pile не дала явного выигрыша в качестве языкового моделирования, что согласуется с работой по GPT-NeoX-20B, но расходится с большинством статей; авторы предполагают, что намеренный upsampling некоторых подмножеств в Pile ведёт себя иначе, чем обычные дубликаты. Параллельное внимание никак не ухудшило качество малых моделей. А "проклятие мультиязычности" в BLOOM проявилось непоследовательно: модель отставала на LAMBADA, PIQA и WSC, но не на WinoGrande, ARC, SciQ или LogiQA.

Практические исследования

Каждое из трёх исследований подобрано так, чтобы задействовать возможности, которые даёт только Pythia.

Исследование гендерной предвзятости заключается в дообучении финального отрезка на изменённых данных. Для дедуплицированных моделей 70M, 410M, 1.4B и 6.9B авторы берут checkpoint и состояние оптимизатора за 21B токенов (7%) до конца обучения и возобновляют процесс на абсолютно тех же данных, но с заменой местоимений мужского рода на женские; для модели 1.4B также обучили версию с заменой на отрезке в 63B токенов (21%). Поскольку порядок и содержание данных во всём остальном идентичны, единственной переменной остаётся частотность местоимений. Стереотипная точность на WinoBias снижается при любом вмешательстве и для всех размеров, а модель 6.9B меняет смещение с про-стереотипного на анти-стереотипное. Предвзятость по CrowS-Pairs также падает - сильнее всего у более крупных моделей, которые изначально были наиболее предвзятыми. Точность на LAMBADA практически не меняется.

Исследование запоминания проверяет, запоминаются ли последовательности, встретившиеся ближе к концу обучения, сильнее. Если использовать определение Карлини и соавторов (Carlini et al.) с 32-токенными промптами и 32-токенными продолжениями для каждой обучающей последовательности, ответ - нет. Количество запомненных последовательностей на батч для модели 12B отлично укладывается в пуассоновский точечный процесс, то есть запоминание распределено равномерно по всему ходу обучения. Отсюда следует практический вывод: размещение чувствительных данных в начале или конце датасета не снизит вероятность запоминания, хотя их размещение в начале хотя бы позволяет выявить проблему на промежуточном checkpoint'е.

Исследование частотности терминов подсчитывает для каждого checkpoint'а, как часто каждый арифметический операнд и каждая сущность из TriviaQA встречались в данных, которые этот checkpoint уже видел. У моделей от 2.8B параметров корреляция между точностью выполнения задачи и частотностью терминов формируется примерно после 65 000 шагов (45% пути обучения); до этого момента, а также у менее крупных моделей, корреляция практически отсутствует. Модели меньше 1B почти никогда не справляются с этими задачами вовсе, даже при наличии 16 примеров в промпте. В задаче умножения разрыв между самыми частыми и самыми редкими операндами увеличивается по мере продолжения обучения.

Место проекта

Pythia определяет "полную открытость" как воспроизводимость вплоть до конкретного батча - это куда более строгий стандарт, чем заявляло большинство последующих релизов. В статье по OLMo авторы называют Pythia и BLOOM самыми открытыми моделями до них и сравнивают собственные checkpoint'ы с Pythia-6.9B, что проводит прямую линию от этого отчёта к olmo-accelerating-the-science-of-language-models. Чем Pythia и не пыталась быть, так это конкурентоспособной по качеству: бюджет в 300B токенов соответствовал GPT-3, тогда как более поздние открытые модели обучались на объёмах вдесятеро больше и выше. Её долговременная ценность - в качестве инструмента для исследований интерпретируемости, запоминания и атрибуции данных, где фиксированный порядок и высокая плотность checkpoint'ов важнее места в рейтингах бенчмарков.