Построим GPT: с нуля и шаг за шагом

Скорее всего все читатели Хабра слышали про chatGPT и знают как он работает. В этой статье, я повторю GPT и обучу на небольшом датасете, в основном пользуясь материалами из видео Andrej Karpathy. Спасибо за помощь в написании статьи и понимании того, что происходит команде из Школы Высшей Математики - у них одна из самых сильных команд по NLP и LLM с преподавателями из ВШЭ и ШАД.

Attention is all you need

Когда мы вводим запрос в chatGPT - ответ следует токен за токеном, слева направо. Это связано с тем, что GPT это sequence to sequence модель, которая расшифровывает как Generative Pre-trained Transformer. Именно эту архитектуру, я и реализую в этой статье.

Transformer исходит из знаменитой статьи 2017 года - Attention is all you need. Изначально, статья была для машинного перевода, однако в последующие несколько лет ее использовали для разного вида задач. В фундаменте chatGPT также лежит архитектура трансфомера.

Датасет

Чтобы повторить chatGPT нужно собрать огромный массив данных (весь интернет) и обучить на большом количестве видеокарт. Так как у меня нет возможности процессить эти данные - я повторю GPT на небольшом датасете всех произведений Шекспира.

Cкачиваем датасет Шекспира

!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt

with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()
print("length of dataset in characters: ", len(text))

В датасете содержится около 1 миллиона (1115394) символов. Если посмотреть на количество уникальных символов - их 65, это различные символы пунктуации и буквы.

!$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz

Это все типы символов, который модель сможет увидеть и предсказать. Например, символа < нет в датасете и модель не сможет его увидеть, и соответственно предсказать.

Токенизация

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

В нашей модели всего 65 различных символов. Переменная stoi переводит символы в индексы, itos индексы символы. Энкодер переводит целую строку в целые числа, а декодер целые числа в строку. Для понимания -- маленькая буква a будет токеном 39, потому что она 39 в списке.

!$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz

# create a mapping from characters to integers
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }
encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers
decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string

print(encode("llm is cool"))
print(decode(encode("llm is cool")))

Это один из самых простых видов токенизации - openAI например использует BPE (Byte pair encoding) токенайзер. Но для упрощения и понимания, мы оставим токенайзер на уровне символов.

Теперь мы можем приступить к токенизации всего датасета. Мы возьмем весь наш датасет из Шекспира, и сделаем энкодинг 1 миллиона символов в целые числа.

import torch # we use PyTorch: https://pytorch.org
data = torch.tensor(encode(text), dtype=torch.long)

Получился тензор размером torch.Size([1115394]) torch.int64

Первые 10 символов выглядят как

tensor([18, 47, 56, 57, 58,  1, 15, 47, 58, 47])

Что соответствует первым 10 символам из текста.

First Citi

F находится на 18 месте, i на 47, r на 56 и так далее

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

n = int(0.9*len(data)) # first 90% will be train, rest val
train_data = data[:n]
val_data = data[n:]

Теперь возьмем блок данных, состоящий из 8 символов.

block_size = 8
print(train_data[:block_size+1])
print(text[:block_size+1])

На выходе это будет в целых числах tensor([18, 47, 56, 57, 58, 1, 15, 47, 58]) и текстом First Cit

Даталоадер

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

torch.manual_seed(42)
batch_size = 4 # how many independent sequences will we process in parallel?
block_size = 8 # what is the maximum context length for predictions?

def get_batch(split):
    # generate a small batch of data of inputs x and targets y
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    return x, y

xb, yb = get_batch('train')

• batch_size = 4: Определяет, сколько независимых последовательностей данных будет обрабатываться параллельно в одном батче.

• block_size = 8: Определяет максимальную длину контекста для предсказания следующего элемента.

Функция get_batch(split):

• Эта функция принимает на вход split (либо ‘train’, либо ‘val’) для выбора набора данных: тренировочный (train_data) или валидационный (val_data).

• Сначала случайным образом выбираются индексы в пределах доступных данных (в зависимости от длины данных и block_size).

• После этого создаются два тензора: x и y:

• x: это входные последовательности длиной block_size.

• y: это целевые последовательности, которые на один шаг смещены относительно x (то есть y содержит следующие элементы после каждого элемента в x).

• Эта функция возвращает батч данных x и соответствующие им целевые значения y.

Если посмотреть на тензор xb, то он выглядит так

tensor([[24, 43, 58,  5, 57,  1, 46, 43],
        [44, 53, 56,  1, 58, 46, 39, 58],
        [52, 58,  1, 58, 46, 39, 58,  1],
        [25, 17, 27, 10,  0, 21,  1, 54]])

Или текстом

Let's he
for that
nt that 
MEO:
I p

Bigram Language Model

Для создания бейзлайн модели будем использовать простую модель - Bigram Language Model.

Биграммная языковая модель оценивает вероятность каждого символа в тексте на основе только одного предыдущего символа, образуя пары символов (биграммы) для предсказания следующего символа в последовательности.

Ее основное преимущество - простая реализация, но она ограничена коротким контекстом, что снижает ее точность при работе с длинными зависимостями в языке.

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

import torch
import torch.nn as nn
from torch.nn import functional as F
torch.manual_seed(42)

class BigramLanguageModel(nn.Module):

    def __init__(self, vocab_size):
        super().__init__()
        
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):

        logits = self.token_embedding_table(idx)
        if targets is None:
          loss = None
        else:
          B, T, C = logits.shape # it will be [4, 8, 65] in our case
          logits = logits.view(B*T, C) # [32, 65]
          targets = targets.view(B*T) # [32]
          loss = F.cross_entropy(logits, targets) 

        return logits, loss
    def generate(self, idx, max_new_tokens):
      for _ in range(max_new_tokens):
        logits, loss = self(idx)
        logits = logits[:, -1, :] # берем последний степ для bigram предсказания
        probs = F.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples = 1)
        idx = torch.cat((idx, idx_next), dim = 1)
      
      return idx


m = BigramLanguageModel(vocab_size)
logits, loss = m(idx=xb, targets=yb)
print(logits.shape, logits, loss)
idx = torch.zeros((1,1), dtype = torch.long)
# принт ниже напринтит чушь, потому что модель не обучена
print(decode(m.generate(idx, max_new_tokens=100)[0].tolist())) 

Создаем оптимайзер для обучения

optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)

И далее мы создаем обучаем модель на 10000 шагах

batch_size = 32
for steps in range(10000): # increase number of steps for good results...

    # sample a batch of data
    xb, yb = get_batch('train')

    # evaluate the loss
    logits, loss = m(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
print(loss.item())

Посмотрим на полученный результат

Hin, ha, hainds t at fl ff f vonearlalininllert's nodo ndere me, Hatheisher: So ineayo in's I dsane

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

Attention

В реализации биграмной модели, токены не были связаны между друг другом. По одному токену - мы предсказывали следующий. Сейчас мы хотим, чтобы токен находящийся например на 5 месте, учитывал все предыдущие токены с 0 по 4 место.

Какой самый легкий способ создать контекст для предыдущих токенов? Пожалуй усреднение всех предыдущих токенов.

# We want x[b,t] = mean_{i<=t} x[b,i]
xbow = torch.zeros((B,T,C)) # нули
for b in range(B): # итерируем по батчам
    for t in range(T): # итерируем по времени
        xprev = x[b,:t+1] # (t,C) 
        xbow[b,t] = torch.mean(xprev, 0)

Для более эффективного вычисления мы используем матричное перемножение. Сумма элементов в матрице wei будет 1.0 и соответственно перемножая wei на x мы получим усреднение.

torch.manual_seed(42)
B, T, C = 4, 8, 2
x = torch.randn(B, T, C)

wei = torch.tril(torch.ones(T, T))
wei = wei / wei.sum(1, keepdim=True)
xbow2 = wei @ x # (B, T, T) @ (B, T, C) ----> (B, T, C)
torch.allclose(xbow, xbow2)

Сейчас если мы посмотрим на тензор Wei - мы увидим, что они распределены у нас равномерно.

tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.5000, 0.5000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.3333, 0.3333, 0.3333, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.2500, 0.2500, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.2000, 0.2000, 0.2000, 0.2000, 0.2000, 0.0000, 0.0000, 0.0000],
        [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.0000, 0.0000],
        [0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.0000],
        [0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250]])

Однако предложения в языке устроены сложнее и имеют неравномерную структуру. Для каждого символа (токена) разные токены имеют различную значимость, в то время как при усреднении все они рассматриваются одинаково. Эту проблему решает механизм self-attention, при котором для каждого токена вычисляются три параметра: key, query и value.

• Key — это то, что содержит данный токен, его содержимое.

• Query — это то, что данный токен “ищет”, или на что он обращает внимание.

• Value — это то, что модель передаст в следующий слой после вычисления внимания.

Self-attention можно рассматривать как механизм коммуникации между токенами. Это похоже на граф, где каждый узел “смотрит” на другие и агрегирует информацию с помощью взвешенной суммы всех узлов, которые на него указывают. Эти веса зависят от данных и определяются на основе контекста токенов.

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

# Устанавливаем фиксированное значение для генератора случайных чисел, чтобы результаты были воспроизводимыми
torch.manual_seed(42)

# Задаем размеры входных данных: 
# B - размер батча (количество примеров в одной итерации),
# T - временные шаги,
# C - количество каналов (размер признакового пространства или число признаков).
B, T, C = 4, 8, 32  # batch, time, channels

# Создаем случайные данные с нормальным распределением, имитируя входные последовательности размером BxTxC
x = torch.randn(B, T, C)

# Определяем размер головы внимания (self-attention). Это размер вектора, на который будем проецировать ключи, запросы и значения.
head_size = 16

# Создаем три линейных слоя для ключей (key), запросов (query) и значений (value).
# Они преобразуют входные данные (размер C) в векторы размером head_size.
key = nn.Linear(C, head_size, bias=False)
query = nn.Linear(C, head_size, bias=False)
value = nn.Linear(C, head_size, bias=False)

# Применяем линейное преобразование к входным данным, получая ключи и запросы для каждого временного шага.
k = key(x)   # (B, T, 16) - матрица ключей
q = query(x) # (B, T, 16) - матрица запросов

# Вычисляем матрицу весов внимания через матричное произведение запросов и транспонированных ключей.
# Это вычисляет сходство между запросами и ключами. Размерность: (B, T, T)
wei = q @ k.transpose(-2, -1)

# Создаем нижнетреугольную матрицу (T x T), которая используется для маскирования будущих временных шагов
# в последовательности (например, для авторегрессии), чтобы модель не "смотрела" вперед во времени.
tril = torch.tril(torch.ones(T, T))

# Применяем маскирование: все элементы выше главной диагонали заменяются на -inf, чтобы исключить их при вычислении softmax.
wei = wei.masked_fill(tril == 0, float('-inf'))

# Применяем softmax к весам, чтобы они стали вероятностями. Это делается по последнему измерению (T).
wei = F.softmax(wei, dim=-1)

# Применяем линейное преобразование к входным данным для получения значений.
v = value(x)

# Вычисляем выходной результат как взвешенное среднее значений с использованием весов внимания.
out = wei @ v  # Размер выхода: (B, T, 16)

Посмотрим на полученную матрицу для первого батча wei[0]. Теперь мы видим, что она у нас не равномерная, и каждый элемент строки показывает степень внимания одного временного шага к другим.

tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.1905, 0.8095, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.3742, 0.0568, 0.5690, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.1288, 0.3380, 0.1376, 0.3956, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.4311, 0.0841, 0.0582, 0.3049, 0.1217, 0.0000, 0.0000, 0.0000],
        [0.0537, 0.3205, 0.0694, 0.2404, 0.2568, 0.0592, 0.0000, 0.0000],
        [0.3396, 0.0149, 0.5165, 0.0180, 0.0658, 0.0080, 0.0373, 0.0000],
        [0.0165, 0.0375, 0.0144, 0.1120, 0.0332, 0.4069, 0.3136, 0.0660]],
       grad_fn=<SelectBackward0>)

Значение 0.3136 в последнем ряду матрицы означает, что для последнего временного шага (восьмого) модель уделяет 31.36% своего внимания седьмому временному шагу при вычислении взвешенного значения (value). Это означает, что седьмой шаг оказывает значительное влияние на вычисление результата для восьмого шага.

Теперь, мы можем создать класс Head для механизма self-attention.

class Head(nn.Module):
    """ one head of self-attention """

    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B,T,C = x.shape
        k = self.key(x)   # (B,T,C)
        q = self.query(x) # (B,T,C)
        # compute attention scores ("affinities")
        wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
        wei = F.softmax(wei, dim=-1) # (B, T, T)
        wei = self.dropout(wei)
        # perform the weighted aggregation of the values
        v = self.value(x) # (B,T,C)
        out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)
        return out

Мы реализовали механизм внимания с использованием скалярного произведения (dot-product attention). Если продолжить чтение статьи, то далее рассматривается Multi-Head Attention, который улучшает способность модели извлекать информацию из разных подпространств признаков, применяя несколько attention-голов параллельно.

class MultiHeadAttention(nn.Module):
    """ multiple heads of self-attention in parallel """

    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.dropout(self.proj(out))
        return out

Обучение модели

Для финального обучения модели я добавил класс, описанный в статье — Feed Forward. Его смысл в том, что он позволяет добавить нелинейность в модель, увеличивая размерность скрытого слоя и обрабатывая данные через два линейных слоя с функцией активации ReLU. Это помогает модели лучше улавливать сложные зависимости в данных, а Dropout используется для предотвращения переобучения.

model = BigramLanguageModel()
m = model.to(device)
# print the number of parameters in the model
print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')

# create a PyTorch optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

for iter in range(max_iters):

    # every once in a while evaluate the loss on train and val sets
    if iter % eval_interval == 0 or iter == max_iters - 1:
        losses = estimate_loss()
        print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

    # sample a batch of data
    xb, yb = get_batch('train')

    # evaluate the loss
    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

Обучение на моем компьютере заняло более часа, но все зависит от типа вашего GPU. Loss уже значительно меньше, чем в обычной Bigram модели - 1.8 вместо 4.7.

step 3900: train loss 1.7219, val loss 1.8732
step 4000: train loss 1.7146, val loss 1.8576
step 4100: train loss 1.7136, val loss 1.8720
step 4200: train loss 1.7060, val loss 1.8653
step 4300: train loss 1.7032, val loss 1.8499
step 4400: train loss 1.7057, val loss 1.8656
step 4500: train loss 1.6907, val loss 1.8477
step 4600: train loss 1.6878, val loss 1.8371
step 4700: train loss 1.6808, val loss 1.8415
step 4800: train loss 1.6689, val loss 1.8457
step 4900: train loss 1.6716, val loss 1.8415
step 4999: train loss 1.6658, val loss 1.8275

Результат

При попытке сгенерировать следующие токены, мы получаем похожий на реального Шекспира текст.

Результат сильно превосходит обычную биграмную модель.

Hin, ha, hainds t at fl ff f vonearlalininllert's nodo ndere me, Hatheisher: So ineayo in's I dsane

Заключение

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