Обучение небольшого GPT на индийских народных сказках: токенизатор, маска и функция потерь
Генерация GPT с 9M параметрами с нуля: почему токенизатор GPT-2 не справляется с словами из народных сказок, как работают BPE и каузальная маска, и почему одна пачка данных дает 4,096 потерь для следующего токена.
Источники: Building LLM, Part 1 — Training: Predicting the Next Token
Небольшой GPT, обученный на индийских народных сказках
Каждый, кто уже создавал модели типа Vision Transformer, знает, что примерно 70 процентов составляют основные компоненты любой языковой модели. Эмбеддинги, механизм внимания, стек остаточных связей и цикл обучения одинаковы в обоих типах моделей. Остается лишь небольшая доработка, и в этом руководстве мы рассмотрим её, создав небольшой GPT с нуля и обучив его на одной видеокарте ноутбука. Корпус данных состоит из индийских народных сказок: историй Джатака, басен в стиле Панчатантра, а также сказок из Пенджаба и Бенгалии.
Серьезная проблема возникает с самого начала. Стандартный токенизатор GPT-2 разбивает такие слова, как Brahman и Rajah, на фрагменты, поскольку ранее он с ними не сталкивался, поэтому вместо этого используется токенизатор, обученный на корпусе данных. В итоге получается модель с 9 миллионами параметров, которая перешла от случайного шума к обученным весам, и любой может воспроизвести эти результаты, запустив обучение, занимающее примерно десять минут.
Почему сначала появилось зрение
Работа с изображениями до текста может показаться неправильным порядком, ведь чат-боты работают с словами. Однако трансформеру безразлично, являются ли его входными данными фрагменты изображений или части предложений. Головы внимания, LayerNorm, мостовые соединения, MLP и цикл оптимизации остаются прежними. Лучший вопрос заключается в том, где этот общий механизм проще всего освоить, и изображения выигрывают по четырем причинам.
- Токен — это просто фрагмент изображения, вырезанный из картинки по фиксированной сетке; поэтому не требуется словарь, токенизатор и решение проблемы неизвестных слов.
- Не существует маски причинности. ViT видит всё изображение сразу, поэтому обычная атенционная модель работает раньше, чем вступают в силу дополнительные правила языка, запрещающие смотреть в будущее.
- Поведение модели видимо: фрагменты изображения, карты атенционных весов и области, на которые ориентируется модель, можно отобразить визуально, причем пиксели кажутся более осязаемыми по сравнению с текстом.
- Обучение начинается с нуля на одной GPU за несколько минут (на данных CIFAR-10); таким образом, каждая цифра — это результат собственных вычислений, а не что-то принятое на веру.
В итоге, как только построен ViT, уже существует полноценный трансформер, просто адаптированный для обработки изображений. Языковая модель добавляет лишь механизм преобразования текста в токены (поиск в таблице, а не вычисления), каузальную маску и прогноз следующего токена. Самая сложная и абстрактная часть — сам язык — становится гораздо проще для понимания, как только машина уже знакома с ним.
Обучение — это просто гадание на следующий токен
Во время обучения всё, что делает языковая модель, сводится к прочтению уже имеющихся токенов и прогнозированию следующего из них.
На рисунке показано прогнозирование следующего токена на основе предыдущих токенов. Если ввести начало предложения «В одной деревне жил бедный брахман», система выдаст вероятность каждого возможного следующего токена. Точка, запятая или слова вроде кто или именованный будут грамматически корректными вариантами. Никто не должен ничего аннотировать: правильным ответом будет тот токен, который действительно следовал в тексте. Именно поэтому один только необработанный текст может обучать языковую модель, и он составляет всю сигнализацию для обучения. Все остальные компоненты существуют для поддержки этого единственного прогноза, причем ViT уже обеспечивает большинство из них.
В коде эта идея реализуется за две строки: сначала запускается модель, затем её оценки нормализуются в распределение вероятностей следующего токена.
logits, _ = model(ids) # ids = the tokens so far, shape (1, T)
probs = F.softmax(logits[0, -1], dim=-1) # one probability per vocabulary token
Остальная часть посвящена тому, как создается и обучается модель, чтобы эти вероятности оказались достаточно хорошими.
Почему важен токенизатор
Модель ViT разделяет изображение на патчи, в то время как языковая модель разделяет текст на токены. Патч вычисляется на основе значений пикселей, а текстовый токен берется из фиксированного словаря. Кто-то должен заранее создать этот словарь, и именно он определяет, с чем модели будет легко справляться.
В первой попытке был использован r50k_base — готовый токенизатор GPT-2, словарь которого был сформирован на основе текстов западных веб-ресурсов примерно 2019 года. На следующем рисунке показано, как он обрабатывает слова из сказок.
Brahman разделяется на Brah и man, Rajah — на Raj и ah, а Panchatantra распадается на пять частей. Ни один из этих терминов не является редким. Слово Rajah встречается 709 раз в корпусе данных, однако инструмент обработки текста никогда не сталкивался с такими словами достаточно часто в своих учебных данных, чтобы выделить для каждого из них отдельный токен. Фрагментация увеличивает длину последовательностей и заставляет модель собирать образ Brahman заново из Brah и man, причем man также совпадает со словами woman и human. При наличии всего около миллиона токенов в учебных данных это представляет собой серьезное препятствие.
Средством является токенизатор, обученный на этих сказках с использованием кодирования пар байтов (BPE). GPT использует тот же алгоритм, только применённый к более подходящему корпусу данных. Исходным материалом служат двенадцать сборников индийских народных сказок из общественного достояния от Project Gutenberg, включая две книги Джатак, два сборника индийских сказок, «Старые дни Декана», «Народные сказки Бенгалии», «Говорящий дрозд», «Сказки Пенджаба», «Двадцать два гоблина», два сборника Санталов и «Сказки о солнце из Южной Индии», что в сумме составляет примерно 809 000 слов.
На рисунке показаны процессы объединения пар байтов, использованные для обучения токенизатора. BPE начинается с 256 исходных значений байтов, так что ничего не может остаться неизвестным; затем повторяющимися операциями объединяются наиболее часто встречающиеся соседние пары в новые токены. После 7 935 операций объединения над сказками словарь насчитывает 8 192 записи, включая целые слова вроде Rajah, Brahman и jackal.
Один из вариантов дизайна требует четкого пояснения: токенизатор GPT-2 заменяется полностью, а не дополняется.
На рисунке показано пользовательское словарь из 8 192 токенов. Добавление токенов к уже имеющимся у r50k 50 257 токенам — разумный подход, когда необходимо сохранить английские навыки крупной предобученной модели, как это делают проекты вроде OpenHathi. Обучение с нуля на небольшом специализированном корпусе — это другая ситуация: словарь r50k, основанный на западных веб-ресурсах, только занимает место, причем рассматриваемые тексты охватывают лишь 37 процентов этого словаря. Новый словарь из 8 тысяч токенов меньше по размеру, 96,6 процента из них используется, и он адаптирован под конкретные тексты, поэтому замена является наилучшим вариантом.
Обучение токенизатора занимает несколько строк кода с использованием библиотеки HuggingFace tokenizers. Скрипт читает каждую книгу в сыром виде, удаляет заголовки и подзаголовки проекта Project Gutenberg и обучает модель BPE.
import glob
from tokenizers import Tokenizer, models, trainers, pre_tokenizers
from llm_transformer.prep_tales import strip_gutenberg, RAW_DIR
# read every raw .txt book and strip its Project Gutenberg header/footer,
# leaving just the story text - one big string per book
tales_texts = [strip_gutenberg(open(p, encoding="utf-8", errors="ignore").read())
for p in sorted(glob.glob(f"{RAW_DIR}/*.txt"))] # 12 books, ~809k words
tok = Tokenizer(models.BPE())
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) # start from bytes
trainer = trainers.BpeTrainer(vocab_size=8192, special_tokens=["<|endoftext|>"],
initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
tok.train_from_iterator(tales_texts, trainer) # learn the 7,935 merges
tok.save("tales_bpe.json")
Файл вывода с расширением .json представляет собой готовый токенизатор — фиксированный список объединений. Его больше не обучают. Его загружают и используют для преобразования текста в идентификаторы и обратно. Один вызов метода .encode() связывает необработанный текст с нейронной сетью.
tok = Tokenizer.from_file("tales_bpe.json") # reload the trained tokenizer
ids = tok.encode("Once upon a time").ids # -> [412, 987, 15, 733] (integers)
idx = torch.tensor(ids).unsqueeze(0) # (1, 4) — a batch of one sequence
vocab_size = tok.get_vocab_size() # 8192 - sets the NUMBER OF ROWS below
В сеть поступают только эти целые числа, а оригинальные слова для неё больше не существуют. Каждый идентификатор — это просто позиция в таблице, а параметр vocab_size определяет количество записей, зарезервированных для слоя эмбеддингов.
От идентификаторов к векторам
Как только текст преобразуется в список идентификаторов, каждый из них становится вектором. Именно здесь нейросеть типа LLM проще по структуре, чем ViT.
На рисунке показано соответствие индекса токенизатора вектору эмбеддинга. В модели ViT для вычисления токена каждого фрагмента из пикселей используется простое умножение матриц. В нейросети типа LLM токен представляет собой дискретный идентификатор, поэтому вектор — это просто номер строки в выученной таблице, что позволяет осуществить быстрый поиск. Токен 789 возвращает строку номер 789, представляющую собой 256-мерный вектор; добавление позиционного эмбеддинга, как в модели ViT, завершает этап обработки входных данных.
В коде слой nn.Embedding представляет собой таблицу, и при его вызове производится поиск по индексам, а не путем умножения. Его атрибут .weight — это обучаемый тензор размером 8192 на 256, содержащий 2,1 миллиона весов, которые обучаются с использованием метода обратного распространения ошибки; отдельная строка изменяется только в тех шагах, когда встречается соответствующий токен.
self.token_embed = nn.Embedding(vocab_size, n_embed) # the 8,192 x 256 table
# .weight IS a learnable (8192, 256) parameter tensor — 2.1M weights, trained by backprop;
# each id's row is nudged only on steps where that token appears in the batch
...
tok_emb = self.token_embed(idx) # idx (B, T) -> vectors (B, T, 256), one row per id
Всё, что идет дальше — а именно механизм внимания, остаточный поток и блоки MLP — идентично структуре Vision Transformer, за одним исключением.
Маска причинно-следственных связей
В ViT каждый фрагмент может уделять внимание всем остальным. Языковая модель не может этого делать, поскольку её задача — предсказывать следующий элемент. Если бы позиция t могла видеть позицию t+1, она бы просто скопировала ответ. Решением является маскировка процесса внимания, чтобы каждый токен видел только себя и более ранние токены.
Маска на рисунке имеет форму нижнего треугольника. До применения функции softmax все будущие позиции получают значение отрицательного бесконечности, которое функция softmax преобразует в вес равный нулю; таким образом, заблокированные токены не вносят никакого вклада. Именно этот треугольник представляет собой единственное архитектурное отличие между механизмом внимания ViT и языковых моделей. Скалярные произведения, масштабирование с использованием квадратного корня от размера блока, функция softmax и взвешенная сумма значений остаются неизменными.
В коде это просто дополнительная строка сверху обычного механизма внимания. Сначала создается матрица единиц нижнетреугольной формы, затем как обычно вычисляются оценки, а позиции, установленные в ноль, заполняются отрицательным бесконечным значением перед применением функции softmax.
# tril = torch.tril(torch.ones(T, T)) — a lower-triangular matrix of 1s, made once
attn = q @ k.transpose(-2, -1) / math.sqrt(head_size) # scores, (B, T, T)
attn = attn.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # blank out the future
attn = F.softmax(attn, dim=-1) # -inf -> weight 0
out = attn @ v
Одна проходка обучает все длины контекста
Маска одной строки скрывает то, что требует более внимательного рассмотрения. Люди часто представляют себе, что языковая модель работает так, будто генерирует текст: добавляет токен в растущий контекст и заново вычисляет механизм внимания на каждом шаге. Это верно во время инференса, но не во время обучения. Во время обучения предложение полностью известно, поэтому значения Q, K и V для всех T токенов вычисляются сразу, формируя полную матрицу внимания T×T за один шаг.
Эта одна проходка возможна только благодаря каузальной маске. Рассмотрим строку t маскированной матрицы: все элементы к правой от диагонали стороне равны отрицательной бесконечности, поэтому после применения функции softmax они становятся нулями. Вычисление out = attn @ V затем даёт строке t сумму взвешенных значений только первых t векторов значений. Таким образом, все длины контекста обрабатываются одновременно.
На рисунке показаны различные прогоны и форма маски во время обучения. Читайте его строчка за строчкой:
- В первой строке учитывается только токен 1, поэтому её выход предсказывает токен 2 на основе контекста из одного токена.
- В третьей строке учитываются токены с 1 по 3, и предсказывается токен 4 на основе контекста из трех токенов.
- В шестой строке учитываются все шесть токенов, и предсказывается токен 7 на основе полного контекста.
Вот ещё один взгляд на проблему. Видимый контекст каждой строки образует нижнетреугольную матрицу, размер которой увеличивается на единицу с каждым разом: 1×1 для первого токена, затем 2×2 и 3×3. Это вложенные лево-верхние блоки большой матрицы, показанные оранжевым цветом; таким образом, один проход размером 6×6 вмещает все меньшие последовательности, а полоса внизу их разворачивает. Все они генерируются внутри одного оператора attn @ V, причём в Python нет циклов, перебирающих позиции. Это объясняет эффективность обучения трансформеров: последовательность из T токенов соответствует T примерам для обучения, при этом контексты длиной от 1 до T обрабатываются параллельно. Та же треугольная маска, которая формирует эти примеры, также не позволяет какой-либо позиции смотреть вперёд во время генерации; процесс происходит по строкам, а стоимость каждого шага снижается за счёт кэширования предыдущих ключей и значений (KV-кэш).
Реальные числа через механизм внимания
До сих пор каждая ячейка означала лишь «присутствовать» или «не присутствовать». При использовании режима множественных вычислений показываются реальные числа для тех же трёх итераций, взятые из блока 0, головы 0 обученной модели для предложения «Однажды жил умный шакал».
На графике представлены вычисления Q, K и V для каждой итерации. Внутри одной строки веса внимания положительны и суммируются до 1, поэтому результат представляет собой выпуклую смесь векторов значений, что является взвешенным средним.
- В первой итерации учитывается один токен с весом 1,00, следовательно
out1 = V1. Первая позиция копирует своё собственное значение, поскольку пока нет ничего для смешивания. - В третьей итерации результат формируется путём смешивания трёх значений с весами примерно 0,20 для первого токена, 0,06 для второго и 0,74 для третьего. Модель в основном опирается на свой собственный токен lived, но всё же учитывает одну пятую от Once, поэтому это действительно смешение.
Это охватывает всё, что делает механизм внимания: оценки из запросов и ключей задают веса, а эти веса смешивают значения. Маскировка гарантирует, что смешивание никогда не выходит за пределы текущего токена, поэтому каждая строка является справедливым обобщением всего прочитанного до этого момента. Реальные головы имеют 32 измерения значений, и в каждом из шести слоев модель содержит по восемь таких голов. На диаграммах отображаются первые 8 измерений одной головы, хотя тот же механизм применяется во всей модели.
Прогнозы T, потери T
Прогноз следующего токена приносит преимущество, отсутствовавшее у ViT. ViT делал один прогноз на изображение — метку класса. При помощи каузальной маски языковая модель может одновременно угадать следующий токен во всех позициях, так что последовательность предоставляет столько же сигналов для обучения, сколько и токенов в ней.
На графике показаны средние значения T прогнозов и функции потерь для входа из T токенов за каждую итерацию. Каждая позиция получает функцию потерь кросс-энтропии, которая измеряет степень удивления модели реальным следующим токеном. Некоторые прогнозы считаются простыми, например a после lived, с значением 0,58. Другие — сложными, например прилагательное clever с значением 7,17, поскольку могло подойти множество других слов. Среднее значение за шесть прогнозов составляет 3,59, и обратная пропагация снижает это значение, используя тот же механизм правила цепочки, что и ViT.
В коде все потери T возникают в результате одного вызова функции cross_entropy. Сетка позиций формата (B, T) преобразуется в один длинный список, и каждый элемент сравнивается со своим истинным следующим токеном.
logits = self.lm_head(x) # (B, T, vocab)
B, T, C = logits.shape
loss = F.cross_entropy(logits.reshape(B*T, C), # every position ...
targets.reshape(B*T)) # ... vs the token that actually followed
Каждая потеря представляет собой задачу классификации словаря
Важно точно определять, что сравнивает функция cross_entropy, поскольку и прогноз, и целевое значение являются векторами, а не отдельными словами. На каждой позиции модуль lm_head генерирует вектор логитов длиной 8,192 элементов, где каждый элемент соответствует оценке для определенного токена словаря, после чего функция softmax преобразует его в распределение вероятностей по всем токенам. Целевое значение — это вектор один-в-горячем той же длины, содержащий 1 для слова, которое действительно следовало, и 0 во всех остальных местах. Таким образом, прогнозирование следующего токена представляет собой задачу классификации, где элементы словаря выступают в роли классов.
На графике показаны логиты в LLM и ViT. В нем объясняется один тонкий момент: столбцы отсортированы по вероятности лишь для удобства восприятия, тогда как фактический результат — это вектор с фиксированной структурой, где слот i всегда соответствует идентификатору токена i. Слово «son» всегда находится по индексу 498 при каждом прохождении модели, будь то обучение или инференс, а функция softmax изменяет только высоту столбцов, но не их порядок. Это объясняет, почему единственная единица в одноразовом целевом векторе находится по индексу 498, а не на первой позиции: токенизатор один раз присвоил каждому токену его идентификатор в качестве соответствующего слота.
Это именно голова классификатора ViT, которая применяет функцию softmax к оценкам классов и оценивает их с помощью функции кросс-энтропия по отношению к меткам в формате one-hot. Различаются два аспекта: количество классов (8 192 токена словаря против примерно 10 меток изображений) и частота использования (на всех позициях T против одного раза на изображение). При работе с целевыми метками в формате one-hot функция кросс-энтропия сводится к отрицательной логарифмической функции вероятности, присвоенной соответствующему слову; поэтому показатель потерь для каждой позиции выглядит как оценка удивления. Слово «son» с вероятностью 9,8 процента приводит к показателю потерь 2,32, что значительно меньше, чем модель ожидала для такого слова; настоящий сюрприз привел бы к гораздо более высоким показателям.
Как формируется пакет данных
Функция reshape(B*T) подразумевает, что пакет данных содержит B последовательностей, каждая из которых вносит свой вклад в показатели потерь для T позиций. Откуда берутся эти B последовательностей, и имеют ли они одинаковую длину?
Люди часто упускают из виду одну деталь: последовательности дообучения не заполняются и у всех они имеют ровно одинаковую длину контекста. Корпус объединяется в один длинный поток токенов, между документами ставится токен <|endoftext|>, а окна фиксированной длины вырезаются в случайных местах. Все окна берутся из одного непрерывного потока, поэтому ни одно из них не может быть коротким, и токенов заполнения не требуется.
Диаграмма показывает путь от обучающего корпуса к группам данных, а затем к показателю потерь. Единственным специальным токеном в потоке является <|endoftext|> (EOT), который служит разделителем, а не заполнителем, и указывает на место окончания одной народной сказки и начало следующей. Окно может пересекать такую границу и содержать символ EOT посередине, и именно так модель понимает, что сказка закончилась и началась новая. Заполнители используются позже, во время тонкой настройки, когда примеры отличаются по длине, и позиции с заполнителями маскируются из расчёта показателя потерь; это отличается от процесса пакетного предобучения.
Потери — вторая область, в которой LLM отличается от ViT, и это различие является положительным. ViT генерировал по одному прогнозу и по одной потере на изображение. Последовательность длиной 128 токенов создавала 128 потерь, среднее значение которых составляло потери для данного образца (3,78 для первого образца на рисунке), а среднее значение потерь 32 образцов в пакете давало потери пакета 3,92, которые минимизируются с помощью алгоритма обратного распространения ошибки. Благодаря этим двум этапам средневзвешивания один шаг снижает среднее значение 32 умноженных на 128, то есть 4096 потерь для следующего токена, что обеспечивает гораздо более сильный сигнал на каждом шаге, чем одна потеря на изображение.
Наблюдение за снижением потерь
Подключите эти потери к оптимизатору и повторяйте процесс. Это и есть настоящая процедура обучения на корпусе рассказов.
График отображает показатели потерь во время обучения и на валидационной выборке в течение эпох. Значение потерь при обучении начинается с 9,2, что соответствует ln(8192) — потерям при случайных предположениях с использованием словаря из 8192 токенов, а наилучшее значение потерь на валидационной выборке достигает 4,73. Кривые остаются близкими друг к другу, с разницей около 1,0, что указывает на то, что при узкой области применения модель способна к обобщению, а не просто запоминанию. Сама модель имеет небольшие размеры.
На графике показана связь между процессом обучения и количеством параметров, которые можно настроить. При 9,0 миллионах параметров модель обучается примерно за десять минут на GPU ноутбука с объемом памяти 6 ГБ. Чуть более половины этих параметров находятся в шести блоках трансформера, а остальные — в таблице эмбеддингов и выходном блоке. Обучение заключается исключительно в постоянной корректировке этих 9 миллионов чисел, так что вероятности появления следующего токена увеличивают вес того токена, который действительно следовал за предыдущим.
В коде цикл выглядит так, как и ожидается, и имеет ту же структуру, что и у ViT: берутся случайные примеры контекста и окна следующих токенов, выполняется процесс передачи сигнала для получения одного значения средней потери, градиенты сбрасываются до нуля, происходит обратная передача градиентов и шаг обновления параметров.
for step in range(train_steps):
xb, yb = get_batch(train_data) # random (context, next-token) windows
_, loss = model(xb, yb) # forward: T predictions -> one mean loss
opt.zero_grad(); loss.backward() # backprop the loss into all 9M parameters
opt.step() # nudge them downhill
В полной версии скрипта предусмотрена схема изменения скорости обучения по формуле косинуса и процедура оценки.
Что нового и что не изменилось
По сравнению с Vision Transformer в LLM добавлено немного нового: эмбеддинги токенов, функционирующие как таблица поиска, каузальная маска, а также функция потерь для следующего токена на каждой позиции, которая классифицирует элементы по всему словарю. Также стоит отметить практический вывод: токенизатор является частью модели, поэтому использование специализированного токенизатора для конкретной области приносит пользу. Механизмы внимания, позиционные эмбеддинги, поток остатков, LayerNorm, блоки MLP и полный цикл обратного распространения ошибки с оптимизатором остаются без изменений.
Реальные числа позволяют понять внутренний механизм работы: как маска причинно-следственных связей позволяет одному проходу оценить все длины контекста сразу, что генерирует attn @ V (каждая строка вывода объединяет векторы значений, доступные для данного токена), как пакеты данных для обучения формируются путем сжатия вместо заполнения нулями, и как потери для одной позиции, представляющие собой классификацию в рамках словаря, объединяются с потерями для других позиций в одно среднее значение.
Веса уже определены, но обученная модель не выполняет никаких действий до тех пор, пока она не запущена. Далее идет процесс инференса: цикл генерации, кэш KV, обеспечивающий высокую скорость работы, и то, что пишет небольшая модель при получении начала текста.