Привет, Хабр! Меня зовут Владимир, и это вторая часть цикла статей по написании и обучению небольшой decoder-only LLM с нуля. В первой части мы вытащили текст, обучили Byte-level BPE токенизатор и собрали pretrain-датасет. Теперь напишем сердце модели - трансформер.
Содержание цикла
Трансформер (вы находитесь здесь)
Сборка и обучение LLM
SFT этап (дообучение на Вопрос-Ответ)
Интеграция с Hugging Face
Содержание может меняться и дополняться ссылками по мере написания
Из чего вообще состоит LLM
Слой эмбеддингов, который преобразует токен в вектор.
Несколько одинаковых блоков трансформера - основная “думающая” часть.
Выходного слоя, который преобразует размышления трансформеров обратно в словарь токенов.
В данной статье и будет рассмотрено создание блока трансформера.
Что такое трансформер
Трансформер - это архитектура нейросети из статьи Google “Attention Is All You Need”. Главная идея: не проталкивать последовательность через RNN/LSTM по одному токену, а обработать все позиции параллельно, связав их механизмом attention (внимания).
┌─────────────────────────────────────────────────────────────────────────────────────────┐ │ Игорь, вставь сюда эту заезженную картинку из статьи гугловской, плиз! │ │ И ссылку `https://arxiv.org/html/1706.03762` обязательно кинь в описание, типа шарим │ └─────────────────────────────────────────────────────────────────────────────────────────┘
Немного научной духоты
Хоть в названии статьи Гугла и написано про “внимание”, сам механизм внимания предложил Дмитрий Богданов (Дзмітрый Багданаў на родном белорусском) в своей статье про перевод с помощью RNN. Внимание по Богданову добавлялось к encoder-decoder RNN и позволял декодеру на каждом шаге обращаться к разным частям входной последовательности. В статье “Attention Is All You Need” рассматривается же механизм Self-Attention: не “декодер смотрит на энкодер”, а “каждый токен смотрит на всех в своём тексте”.
В оригинальном трансформере было две башни - энкодер и декодер. Энкодер “видит” все токены последовательности, декодер - только те, что левее него (поэтому в декодере не просто self-attention, а masked self-attention). На оригинальной архитектуре построена модель T5 от Google. Encoder-only архитектуры легли в основу BERT-подобных моделей, оставляем только декодер - и получаем наши любимые GPT, LLaMA и им подобные.
Ладно, освежили в памяти, что такое трансформер, переходим к его ядру - Multi-Head Attention
Multi-Head Attention
Вот он, родимый:
В основе механизма лежат три матрицы - W_q (Query), W_k (Key) и W_v (Value). Каждый токен во входной последовательности преобразуется в три вектора с помощью обучаемых линейных проекций:
Query (Q) — вектор, который отвечает на вопрос: “Что я ищу в других элементах последовательности?”. Он описывает, какой контекст нужен текущему токену.
Key (K) — вектор, который отражает свойства других токенов. С его помощью модель сравнивает текущий запрос с характеристиками остальных элементов.
Value (V) — вектор, который содержит информацию, которую модель хочет получить, если совпадение окажется значимым.
Всё ясно и понятно, правда?
А вот мне лично до сих пор эта аналогия мозг взрывает, хотя я как-то пару вечеров с разными LLM на эту тему плотно общался.
Для себя я выбрал другую аналогию. На первом этапе (там где X·W_q и X·W_k перемножаются) мы вычисляем сходство (или близость) токенов между собой. Если бы матриц W_q и W_k не было, то сколько раз ни умножай X @ Xᵀ - результат один и тот же. Для этого и нужны матрицы W_q и W_k - они переносят исходные эмбеддинги из общего эмбеддинг-пространства в два отдельных, обучаемых подпространства. Типа как в CV, где каждый свёрточный слой в итоге начинает выделять какие-то свои признаки (вертикальные линии, пятачки, хвостики-крючки), так и в MHA - разные головы могут фокусироваться на разных аспектах входа (синтаксис, семантика, зависимости и т.п.).
Выход первого этапа строит матрицу отношений между токенами (aka Матрица Внимания): кто с кем связан и насколько сильно. На втором этапе в игру входит X·W_v: мы берём содержимое токенов, тоже спроецированное в своё обучаемое подпространство, и смешиваем его согласно найденным весам внимания. В аналогии с тем же CV: сначала решили, на какие признаки смотреть, а потом через V протащили дальше свой тип сигнала — где-то цвет, где-то контраст, где-то контуры.
Перейдём к реализации класса MultiHeadAttention:
class MultiHeadAttention(nn.Module): def __init__(self, emb_dim: int, n_heads: int, context_length: int, dropout: float = 0.1, qkv_bias: bool = False): super().__init__() if emb_dim % n_heads != 0: raise ValueError('emb_dim должно быть кратно n_heads') self.model_dim = emb_dim self.n_heads = n_heads self.head_dim = self.model_dim // n_heads self.qkv_proj = nn.Linear(self.model_dim, 3 * self.model_dim, bias=qkv_bias) self.out_proj = nn.Linear(self.model_dim, self.model_dim, bias=qkv_bias) self.dropout = nn.Dropout(dropout) self.register_buffer('mask', torch.tril(torch.ones(1, 1, context_length, context_length)))
Головы многоголового внимания формируются делением входного эмбеддинга на равные части, поэтому первым делом мы проверяем, что можем поделить размерность вектора на число голов без остатка. Далее формируем матрицы W_q, W_k и W_v. Это должно быть три раза по количеству голов матриц, каждая размером emb_dim // n_heads. Но тогда при каждом вызове приходилось бы выполнять 3·n_heads отдельных умножений мини-матриц, а это лишние чтения памяти, лишние запуски CUDA-ядер, лишние накладные расходы. Поэтому вместо отдельных мини-матриц делают одну большую, размерностью (emb_dim, 3·emb_dim) - умножений столько же, но за счёт сокращения накладных расходов производительность вырастает. После QKV у нас, согласно схеме, идёт выходной линейный слой. Также добавим слой Dropout, который при обучении случайно выключает часть связей между токенами, чтобы модель не переобучилась на фиксированные паттерны.
Последний оператор нужен для создания маски:
torch.onesсоздаёт матрицу из единичекtorch.trilзаполняет верхний треугольник матрицы нулями (относительно диагонали [0, 0] - [N, N])register_bufferзапоминает полученную матрицу под именемmask, привязывая её к модели - теперь на матрицу иmodel.to(device)корректно отработает, и вmodel.state_dict()она по умолчанию попадёт.
Теперь метод forward
class MultiHeadAttention(nn.Module): # Предыдущий код def forward(self, x: torch.Tensor) -> torch.Tensor: # Шаг 1 batch_size, n_tokens, emb_dim = x.size() # Шаг 2 qkv = self.qkv_proj(x) q, k, v = qkv.split(self.model_dim, dim=-1) # Шаг 3 queries = q.view(batch_size, n_tokens, self.n_heads, self.head_dim).transpose(1, 2) keys = k.view(batch_size, n_tokens, self.n_heads, self.head_dim).transpose(1, 2) values = v.view(batch_size, n_tokens, self.n_heads, self.head_dim).transpose(1, 2) # Шаг 4 attn_weights = (queries @ keys.transpose(-2, -1)) / (self.head_dim ** 0.5) # Шаг 5 attn_weights = attn_weights.masked_fill(self.mask[:, :, :n_tokens, :n_tokens] == 0, float('-inf')) attn_weights = self.dropout(torch.softmax(attn_weights, dim=-1)) # Шаг 6 out_per_head = attn_weights @ values # Шаг 7 out = out_per_head.transpose(1, 2).contiguous().view(batch_size, n_tokens, emb_dim) return self.out_proj(out)
Объясню по шагам, что происходит в методе.
На вход метода подаётся батч векторизированных токенов, следовательно, имеем тензор размерностью [Размер батча, Кол-во токенов, Размерность слоя эмбеддингов]. Все три размера запоминаем для дальнейших преобразований (Шаг 1).
На шаге 2 мы вычисляем наши Q, K и V - сначала прогоняем входной тензор через объединённый слой, а затем делим на три матрицы по последней размерности (которая 3 * emb_dim).
Следующим шагом мы должны поделить наши матрицы на головы и подготовить их для дальнейших вычислений. Для этого мы, с помощью метода .view преобразуем последнюю размерность (которая сейчас emb_dim) в матрицы размерностями (n_heads, head_dim), а методом .transpose меняем местами порядок измерений, чтобы получить в итоге матрицу с размерностями (batch_size, n_heads, n_tokens, head_dim) для дальнейшего “поголовного” перемножения. Повторяем со всеми тремя матрицами.
На четвёртом шаге вычисляем близость токенов, для чего надо скалярно умножить матрицу Q на транспонированную матрицу K. При этом при транспонировании первые две размерности (батч и количество голов) размерности не трогаем. Полученную матрицу, размера (batch_size, n_heads, n_tokens, n_tokens), нормируем, поделив на корень из размерности одной головы.
На пятом шаге мы маскируем полученные значения сходства (заменяем минус бесконечностью все значения верхнего треугольника), затем softmax, который из сырых близостей делает нормированные веса внимания, и dropout.
На шестом шаге мы перемножаем веса внимания с матрицей V, получая опять матрицу размера (batch_size, n_heads, n_tokens, head_dim)
Седьмой шаг - это слияние голов. В нём мы выполняем преобразования, обратные шагу 3. Сначала возвращаем на “измерение” токенов, затем, с помощью contiguous() мы уплотняем вектор (после этой операции весь вектор неразрывно лежит в памяти), после чего сливаем измерения (n_heads, head_dim) в одну emb_dim.
Возвращает функция значения седьмого шага, пропущенные через выходной линейный слой.
Осталось реализовать вторую половину трансформера - Feed Forward слой
Feed Forward Network
По канону, данная подсеть должна состоять из двух линейных преобразований с активацией ReLU между ними, при этом внутренняя размерность должна быть в 4 раза больше входа-выхода:
class FeedForward(nn.Module): def __init__(self, emb_dim: int, dropout: float = 0.1): super().__init__() self.network = nn.Sequential( nn.Linear(emb_dim, 4 * emb_dim), nn.GELU(), nn.Linear(4 * emb_dim, emb_dim), nn.Dropout(dropout) ) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.network(x)
На практике GELU на языковых задачах даёт немного лучший loss, чем ReLU, да и в GPT2 была именно GELU. Поэтому и я взял её, а не каноничный ReLU.
Детали готовы, собираем трансформер.
Сборка трансформера
class TransformerBlock(nn.Module): def __init__(self, emb_dim: int, n_heads: int, context_length: int, dropout: float = 0.1, qkv_bias: bool = False): super().__init__() self.attn = MultiHeadAttention( emb_dim=emb_dim, n_heads=n_heads, context_length=context_length, dropout=dropout, qkv_bias=qkv_bias) self.ffn = FeedForward(emb_dim=emb_dim, dropout=dropout) self.ln_1 = nn.LayerNorm(emb_dim) self.ln_2 = nn.LayerNorm(emb_dim) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.attn(self.ln_1(x)) x = x + self.ffn(self.ln_2(x)) return x
Вот, собственно, и весь трансформер. В конструкторе класса создаём слои MultiHeadAttention и FeedForward, а также для слоя нормализации. В методе прямого прохода реализованы residual connections (ну или skip connections), заимствованные из CV архитектуры ResNet. Смысл в том, что после некоторого количества слоёв мы прибавляем к выходу слоёв входной тензор. Это позволяет градиентам свободно течь через сеть при обратном распространении ошибки.
Самое главное отличие данного трансформера от каноничного из 2017 года - это преднормализация, вместо классической постнормализации. В статье “On Layer Normalization in the Transformer Architecture” описывается, что при использовании Pre-LN можно избежать нескольких серьёзных проблем, имеющихся у Post-LN архитектуры: Pre-LN решает проблему обязательного “прогрева” (warm-up) скорости обучения, градиенты при Pre-LN ведут себя более предсказуемо и стабильно.
Итоги
Поздравляю. Самопальный Трансформер готов. В следующей части мы, на основе нашего блока, соберём небольшую LLM, напишем код для предварительного обучения модели, а также код для генерации текста.
Код проекта тут