Навчання невеликого GPT на індійських народних казках: токенайзер, маска та втрати.

Невеликий GPT, навчений на індійських народних казках

Кожен, хто вже створював Vision Transformer, приблизно знає 70 відсотків того, що потрібно для мовної моделі. Ембеддинги, механізм уваги, стек залишкових значень та цикл навчання — у всіх цих моделях використовується одна й та сама архітектура. Решта — це особливості конкретної моделі, і цей огляд розглядає їх шляхом створення невеликого GPT з основ та його навчання з нуля на GPU одного ноутбука. Корпус даних складається з індійських народних казок: історій з жанру Джатака, байок у стилі Панчатантра та казок з Пенджабу та Бенгалії.

Справжня проблема виникає вже на початку. Стандартний токенайзер GPT-2 розбиває такі слова, як Brahman та Rajah, на фрагменти, оскільки ніколи раніше не стикався з ними, тому замість нього використовується токенайзер, навчений на корпусі даних. У кінцевому підсумку ми отримуємо модель із 9 мільйонів параметрів, яка перейшла від випадкового шуму до навчених ваг, і будь-хто може відтворити ці результати, запустивши тренування, яке триває приблизно десять хвилин.

Чому спочатку з’явилася візуалізація

Розглядати зображення перед текстом може здатися неправильним порядком, адже чат-боти працюють із словами. Проте трансформер не зважає на те, чи є його вхідними даними фрагменти фотографій чи частини речень. Голови уваги, LayerNorm, з’єднання типу skip connection, 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, включаючи дві книги Jataka, два томи індійських казок, «Старі дні Деккану», «Народні казки Бенгалії», «Розмовний дрозд», «Казки Пенджабу», «Двадцять два гобліни», дві збірки казок народу Сантал та «Казки сонця з Південної Індії», що разом становлять близько 809 000 слів.

На графіку показано процес об’єднання пар байтів, який використовувався для навчання токенайзера. BPE починається з 256 первинних значень байтів, тож немає нічого невідомого, після чого постійно об’єднуються найпоширеніші сусідні пари в новий токен. Після 7 935 об’єднань у межах казок словник містить 8 192 елементи, серед яких цілі одиниці, такі як Rajah, Brahman та jackal.

Одне з рішень щодо дизайну потребує чіткого пояснення: токенайзер GPT-2 було повністю замінено, а не розширено.

На малюнку показаний власний словниковий запас, що складається з 8 192 токенів. Додавання токенів до існуючих 50 257 у r50k є розумним підходом, коли потрібно зберегти англійські навички великої попередньо навченої моделі, як це роблять проекти на кшталт OpenHathi. Інша ситуація — навчання з нуля на невеликому спеціалізованому корпусі текстів, де словниковий запас r50k, орієнтований на західні веб-ресурси, лише займає місце, причому оповідання охоплюють лише 37 відсотків цього запасу. Новий словниковий запас у 8 000 токенів є меншим, 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 та LLM. Дотичні добутки, масштабування за допомогою квадратного кореня розміру блоку, функція 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, тож це справжня суміш.

У шостому запуску ваги розподіляються між усіма шістьма значеннями: по 0,26 на кожне з there та jackal. Цей змішаний 32-вимірний вектор підсумовує весь префікс для голови, а наступні шари перетворюють його на прогноз.

Це охоплює все, що робить механізм уваги: оцінки з запитів та ключів визначають ваги, а ваги змішують значення. Маскування гарантує, що це змішування ніколи не виходить за межі поточного токена, тож кожен рядок є справедливим підсумком усього прочитаного досі. Справжні голови мають 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 перетворює його на розподіл ймовірностей між усіма ними. Ціль — це вектор типу one-hot такого ж розміру, у якому 1 знаходиться біля слова, що дійсно слідувало, а 0 — всюди інде. Таким чином, прогнозування наступного токена є задачею класифікації, де слова словника є класами.

На цьому графіку порівнюються логітти в моделі LLM та моделі ViT. Тут пояснюється один важливий момент: смуги відсортовані за ймовірністю лише для кращої читабельності, тоді як фактичний результат — це вектор із фіксованою структурою, де поле i завжди належить до ідентифікатора токена i. Слово „son“ залишається на індексі 498 під час кожного проходження моделі, чи то під час навчання, чи під час інференсу, а функція softmax змінює лише висоти елементів у цих полях, але ніколи їхній порядок. Це пояснює, чому єдиний елемент значення one-hot у форматі 1 знаходиться на індексі 498, а не на першій позиції: токенайзер один раз присвоїв кожному токену його ідентифікатор як номер його поля.

Це саме голова класифікатора ViT, яка застосовує функцію softmax до оцінок класів та оцінює їх за допомогою крос-ентропії відносно однорозмірного міткування. Є дві відмінності: кількість класів (8 192 токени словникового запасу проти приблизно 10 міток для зображень) та частота використання (у всіх позиціях T проти один раз на зображення). При однорозмірному цільовому міткуванні крос-ентропія також зводиться до мінус логарифму ймовірності, призначеної правильному слову, тому втрата для кожної позиції нагадує оцінку здивування. Слово „son“ із ймовірністю 9,8 відсотка коштує 2,32 — це слово, яке модель сильно очікувала побачити з набагато меншою втратою, тоді як справжнє здивування призводило б до значно вищої втрати.

Як створюється пакет даних

Функція reshape(B*T) просто вказує, що пакет містить B послідовностей, кожна з яких внескає свою частку втрат для T позицій. То звідки беруться ці B послідовностей, і чи всі вони однакової довжини?

Люди часто не помічають однієї деталі: послідовності передтренування не доповнюються та всі вони мають точно однакову довжину контексту. Корпус об’єднується в один довгий потік токенів, між документами розміщується токен <|endoftext|>, а вікна фіксованої довжини вирізаються у випадкових місцях. Усі вікна походять з одного безперервного потоку, тому жодне з них не може бути коротким, і додаткові токени для доповнення не потрібні.

Ця діаграма відображає шлях від навчального корпусу до пакетів даних, а потім до показника втрат. Єдиним спеціальним токеном у потоці є <|endoftext|> (EOT), який використовується як розділник, а не для заповнення, і позначає місце закінчення однієї народної казки та початок наступної. Вікно може перетинати таку межу та містити EOT посередині, і саме так модель розуміє, що історія закінчилась та почалась нова. Заповнення з’являється пізніше, під час тонкої налаштування, коли приклади відрізняються за довжиною, і позиції з заповненням маскуються у розрахунку втрат, що є іншим підходом порівняно з пакетним попереднім навчанням.

Втрата є другою областю, де LLM відрізняється від ViT, і ця різниця є позитивною. ViT генерував по одній прогнозації та по одній втраті на кожне зображення. Послідовність довжиною 128 токенів створює 128 втрат, які усереднюються для отримання втрати цього зразка (3,78 для зразка 1 на малюнку), а 32 втрати зразків у пакеті усереднюються для отримання втрати пакета 3,92, яку мінімізує алгоритм backprop. Завдяки цим двом етапам усереднення один крок знижує середнє значення 32 на 128, тобто 4 096 втрат для наступного токену, що забезпечує значно більше сигналу на кожен крок, ніж одна втрата на зображення.

Спостереження за зниженням втрат

Під’єднайте цю втрату до оптимізатора та повторіть процедуру. Це і є справжнє навчання на корпусі оповідань.

Графік показує зниження втрат під час навчання та на валідаційній даних протягом епох. Навчання починається на рівні 9,2, що дорівнює ln(8192) — втратам від випадкових припущень у словнику з 8 192 токенів, а найкращий показник втрат на валідації досягає 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, який забезпечує швидкість, та те, що пише невеликий модель, коли отримує початок історії.