Визуальное руководство по LLM, часть 2: внутри архитектуры Transformer

Оригинал опубликован Rohit Kumar Tiwari в X.

В этой серии из двух частей мы шаг за шагом разбираем ключевые компоненты архитектуры больших языковых моделей (LLM) — так, чтобы это было понятно даже тем, кто только начинает знакомиться с ИИ.

План серии

Часть 1 показывает, как письменный язык разбивается на токены, а затем превращается во входные embedding-векторы, с которыми может работать модель. В ней объясняются токенизация, token embeddings и positional embeddings.

Часть 2 показывает, как эти входные embedding-векторы проходят через transformer block и превращаются в выход модели. Здесь рассматриваются self-attention, causal self-attention, masked multi-head attention, feedforward neural networks, выбор следующего токена и то, как обучение постепенно улучшает модель.

В первой части визуального руководства по LLM мы разобрали, как большие языковые модели берут человеческий язык и превращают его в форму, понятную машине. Мы взяли предложение «Every moment is a beginning», разбили его на токены, присвоили им числовые ID и преобразовали их в плотные векторные представления — token embeddings. Затем добавили positional embeddings, чтобы модель понимала порядок слов.

Во второй части мы рассмотрим Transformer Block и Output Layer.

Ниже — разбор каждого компонента transformer block по отдельности.

Один из самых важных компонентов — Masked Multi-Head Attention. Но прежде чем перейти к нему, нужно понять механизмы Self Attention и Causal Self Attention. После этого гораздо проще разобраться, как работает Masked Multi-Head Attention. Также мы рассмотрим Feedforward Neural Networks, Residual Connections и Layer Normalization, чтобы увидеть, как все элементы transformer block работают вместе.

1. Механизм Self Attention

Язык сильно зависит от контекста. Если прочитать слово «beginning», его точный смысл во многом определяется словами, которые стоят перед ним. Механизм self-attention позволяет модели понять, как слова связаны друг с другом. Вместо того чтобы рассматривать каждое слово изолированно, каждый токен «смотрит» на другие токены и решает, какие из них важнее.

Attention(Q, K, V) = softmax((QK^T) / sqrt(d_k)) · V

В первой части мы уже получили input embeddings после сложения token embeddings и position embeddings.

Пример: Every moment is a beginning

Token IDs: {”Every”: 15745, “moment”: 4205, “is”: 382, “a”: 261, “beginning”: 10526}

Итоговые input embeddings — сумма token embeddings и positional embeddings:

• “15745” (position 1): [-0.5880, 0.3486, 0.6603] + [-0.9178, 0.9045, -2.0975]
• “4205” (position 2): [-0.2196, -0.3792, 0.7671] + [1.1558, -1.2157, 0.1295]
• “382” (position 3): [-1.1925, 0.6984, -1.4097] + [1.0937, 0.2066, 3.1815]
• “261” (position 4): [0.1794, 1.8951, 1.3689] + [0.0967, 1.4086, 0.1915]
• “10526” (position 5): [-1.6033, -1.3250, 0.1784] + [-0.1562, 0.2446, 4.0124]

После поэлементного сложения получаем:

• “15745” (position 1): [-1.5058, 1.2531, -1.4372]
• “4205” (position 2): [0.9362, -1.5949, 0.8966]
• “382” (position 3): [-0.0988, 0.905, 1.7718]
• “261” (position 4): [0.2761, 3.3037, 1.5604]
• “10526” (position 5): [-1.7595, -1.0804, 4.1908]

Чтобы понять слово “beginning” в нашем предложении, модель обращает внимание на слова вроде “moment” и “Every”. Это даёт контекст и помогает уловить идею: каждый момент может быть новым началом.

Шаг 1. Модель использует три вектора, полученные из input embeddings: Queries (Q), Keys (K) и Values (V).

Шаг 2. Считается dot product между всеми Queries и Keys, чтобы измерить, насколько хорошо они совпадают.

Attention Scores = QK^T

Шаг 3. Результат масштабируется на квадратный корень из размерности ключа d_k, чтобы значения оставались стабильными во время обучения.

Scaled Attention Scores = QK^T / sqrt(d_k)

Шаг 4. Применяется softmax, чтобы получить attention weights.

Attention Weights = softmax(QK^T / sqrt(d_k))

Шаг 5. Вычисляются context vectors.

Context Vectors = Attention Weights · V = Attention(Q, K, V)

Полный self-attention: после attention каждый токен содержит информацию, собранную из других токенов последовательности. Это базовая идея трансформеров.

Causal Masking: скрываем будущие токены

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

Механизм causal self-attention использует маску, которая запрещает токенам смотреть на будущие токены. Разрешённый паттерн attention образует нижнетреугольную матрицу: каждый токен может обращаться только к самому себе и к предыдущим токенам.

[1 0 0 0 0]
[1 1 0 0 0]
[1 1 1 0 0]
[1 1 1 1 0]
[1 1 1 1 1]

• 1 означает, что attention разрешён.
• 0 означает, что attention заблокирован.

На практике заблокированные позиции маскируются: их attention scores заменяются на отрицательную бесконечность перед softmax. После softmax эти позиции получают вероятность 0, и модель не может обращать внимание на будущие токены.

Это гарантирует:

• токен 1 видит только себя;
• токен 2 видит токены 1 и 2;
• токен 3 видит токены 1, 2 и 3.

Теперь каждый токен может учитывать только себя и предыдущие токены.

2. Механизм Masked Multi-Head Attention

Одна attention head имеет ограниченную способность улавливать множество разных отношений в языке. Несколько attention heads позволяют модели одновременно учить разные паттерны. Чтобы параллельно учитывать грамматику, смысл, дальние зависимости и связи «подлежащее—объект», модель использует Masked Multi-Head Attention. Несколько attention heads работают одновременно и смотрят на одно и то же предложение по-разному.

Для предложения “Every moment is a beginning” разные heads могут фокусироваться на разных нюансах:

• Attention Head 1 (смысл): связывает “moment” ↔ “beginning”, чтобы уловить идею обновления.
• Attention Head 2 (грамматика): связывает “Every” ↔ “moment”, чтобы понять, что “Every” описывает “moment”.
• Attention Head 3 (структура): связывает “is” ↔ “beginning”, чтобы зафиксировать главное утверждение предложения.

Как это работает:

1. Input embeddings делятся на меньшие части — heads.
2. Каждая head независимо выполняет attention в своём подпространстве.
3. Выходы всех heads объединяются через concatenation.
4. Финальный linear layer объединяет эту разнообразную информацию в единое контекстно-богатое представление.

3. Feedforward Neural Networks (FFN)

Attention позволяет токенам обмениваться информацией по всей последовательности. Но после того как информация собрана, каждый токен всё ещё нуждается в индивидуальной обработке, чтобы модель могла выучить сложные паттерны. Эту роль выполняет feedforward network — часто его называют FFN или MLP block.

Feedforward neural network обычно состоит из двух linear layers с activation function между ними, например GELU. Сеть временно расширяет hidden dimension, чтобы модель могла выучить более сложные зависимости.

1. Linear layer: временно расширяет hidden dimension, часто в 4 раза, давая модели пространство для сложных паттернов.
2. Activation function: добавляет нелинейность, чтобы модель могла понимать сложные, не прямолинейные отношения.
3. Second linear layer: сжимает размерность обратно до исходного размера.

4. Residual Connections

Когда сети становятся глубже, gradients — сигналы, с помощью которых модель обучается, — могут становиться слишком маленькими или слишком большими во время backpropagation. Это известно как vanishing или exploding gradient problem. Когда gradients исчезают, ранние слои обучаются очень медленно, потому что обучающий сигнал ослабевает.

Residual connections, или skip connections, решают эту проблему. Вместо того чтобы учить полностью новое представление, слой учит residual update, который добавляется к исходному входу.

Output = x + Sublayer(x)

Residual connections используются вокруг Masked Multi-Head Attention и FFN layers. Они помогают трансформерам обучать более глубокие сети, стабилизировать gradients и сохранять исходную информацию о словах.

5. Layer Normalization

Когда данные проходят через множество слоёв, activation values могут становиться слишком большими или слишком маленькими, замедляя обучение. Layer Normalization стабилизирует эти значения, нормализуя признаки каждого токена независимо. Это помогает удерживать activations в стабильном диапазоне во время обучения.

Для token embedding:

x = [x_1, x_2, x_3]

LayerNorm считает среднее:

μ = (1 / n) Σ x_i

Дисперсию:

σ² = (1 / n) Σ (x_i - μ)²

И нормализованный выход, где epsilon — маленькое число, добавленное для предотвращения деления на ноль:

x̂_i = (x_i - μ) / sqrt(σ² + ε)

Layer normalization применяется внутри каждого transformer block несколько раз. Благодаря этому обучение становится значительно быстрее и надёжнее.

Transformer Block

Transformer block объединяет основные компоненты архитектуры Transformer. GPT-модели строятся путём последовательного наложения множества таких блоков друг на друга. За счёт этого представления токенов становятся всё более информативными на каждом слое. Современные GPT-модели используют pre-norm design: layer normalization применяется перед attention и feedforward operations.

Поток через transformer block:

1. Layer Normalization
   Нормализует входные представления, повышая стабильность обучения.
2. Masked Multi-Head Attention
   Позволяет каждому токену собирать информацию о себе и предыдущих токенах, не давая доступа к будущим токенам.
3. Residual Connection (Add)
   Исходный вход добавляется обратно к attention output, помогая сохранять информацию и улучшать gradient flow.
4. Layer Normalization
   Повторно нормализует обновлённые представления перед дальнейшей обработкой.
5. Feedforward Network (FFN)
   Применяет нелинейные преобразования, чтобы выучить более сложные паттерны и отношения.
6. Residual Connection (Add)
   Вход перед второй Layer Normalization добавляется к FFN output, сохраняя информацию и одновременно включая новые преобразования.

Примечание: во время обучения после attention и feedforward operations часто применяется dropout. Это помогает уменьшить overfitting и улучшает способность модели к обобщению.

Современные GPT-модели складывают много transformer blocks друг на друга. Каждый блок уточняет представления токенов.

Output Layer

После прохождения через все Transformer blocks модель получает context-aware representation для каждого токена. Она производит logits для каждой позиции токена, но для генерации следующего токена использует logits из последней позиции последовательности — в нашем примере это “beginning”.

Это представление проходит через output layer:

1. Linear Layer (Unembedding)
   Финальный вектор проецируется в вектор размером со словарь модели. Полученные значения называются logits: по одному logit для каждого возможного токена, который модель может сгенерировать.
2. Softmax Function
   Logits преобразуются в probability distribution по всему словарю. Теперь каждое значение отражает оценочную вероятность того, что соответствующий токен станет следующим в последовательности.

Затем модель выбирает следующий токен. Например, если токен “.” имеет самую высокую вероятность, он может быть выбран как предсказание.

Поскольку большие языковые модели генерируют текст по одному токену, предсказанный токен добавляется к исходному входу.

Последовательность:

Every moment is a beginning

становится:

Every moment is a beginning .

Эта обновлённая последовательность снова подаётся в модель, которая повторяет тот же процесс: tokenization, token embeddings, positional embeddings, transformer blocks и output layer, чтобы предсказать следующий токен. Повторяя этот цикл, модель генерирует текст постепенно — один токен за другим.

GitHub Repository

Автор также предлагает поставить звезду и склонировать репозиторий GitHub: проект по созданию ChatGPT-подобной LLM с нуля на PyTorch, с пошаговыми объяснениями.

Код: https://github.com/analyticalrohit/llms-from-scratch

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