Как работает forward pass LLM: от токена до вероятности следующего токена?

Краткий тезис

Forward pass LLM — это последовательность преобразований, превращающих входной текст в распределение вероятностей следующего токена. Процесс включает токенизацию, эмбеддинг, позиционное кодирование, прогон через стопку слоёв трансформера (self-attention + feed-forward с residual-связями и нормализацией), финальную нормализацию и проекцию на словарное пространство через LM head с softmax. Понимание каждого этапа необходимо для диагностики, оптимизации и fine-tuning моделей.

1. Токенизация: текст → последовательность целых чисел

Токенизация — разбиение входного текста на единицы (токены) по словарю модели. Современные LLM используют BPE (Byte-Pair Encoding), WordPiece или SentencePiece (Unigram). Результат — массив целочисленных идентификаторов токенов.

• Пример: "Привет, мир!" → [143, 255, 7, 98, 3] (условно).
• токены|Специальные токены: <bos>, <eos>, <pad>, <unk>.
• Размер словаря vocab_size (типично 32k–128k).

Почему это важно Качество токенизации напрямую влияет на покрытие редких слов, длину последовательности и вычислительную эффективность.

2. Embedding lookup: токены → плотные векторы

Каждый токен преобразуется в вектор размерности d_model (например, 4096 для LLaMA-7B) через embedding matrix размера vocab_size × d_model.

• Вход: (batch_size, seq_len) целых чисел.
• Выход: (batch_size, seq_len, d_model) вещественных чисел.
• Эмбеддинги обучаются совместно с моделью.

Weight tying (связывание весов) — часто матрица эмбеддингов и матрица LM head (см. раздел 6) разделяют одни и те же веса, что уменьшает число параметров.

3. Positional encoding: добавление информации о порядке

Transformer по своей природе инвариантен к порядку токенов, поэтому необходима позиционная кодировка.

• Absolute positional encoding (как в оригинальном Transformer): синусоиды фиксированной частоты.
• Relative positional encoding (RoPE — Rotary Position Embedding): вращение векторов в пространстве эмбеддингов в зависимости от позиции. RoPE используется в LLaMA, Mistral, Qwen и др.
• ALiBi (Attention with Linear Biases): добавление линейного штрафа к attention scores в зависимости от расстояния.

RoPE работает путём применения поворотной матрицы к парам компонент эмбеддинга:

RoPE(x, pos) = x * cos(pos * θ) + rotate_half(x) * sin(pos * θ)

Выход: (batch_size, seq_len, d_model) с внедрённой позиционной информацией.

4. Слои трансформера (×N)

Каждый слой состоит из двух подблоков: Multi-Head Self-Attention и Feed-Forward Network (FFN), с residual-связями и layer normalization (Pre-Norm или Post-Norm).

4.1 Multi-Head Self-Attention

• Вход: X размерности (batch, seq, d_model).
• Линейные проекции на Q, K, Vдля каждого изh голов.
• Для каждой головы: Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V.
• Masked attention (при генерации): маскируем будущие токены, чтобы модель не «подглядывала».
• Конкатенация голов и финальная проекция.

Параметры d_model, h (число голов), d_k = d_model / h.

4.2 Feed-Forward Network (FFN)

• Обычно два линейных слоя с нелинейностью (ReLU, GELU, SwiGLU).
• В LLaMA используется SwiGLU: FFN(x) = (xW_gate ⊙ silu(xW_up)) W_down.
• Расширение скрытой размерности: d_ff (обычно 4× d_model).

4.3 Residual connection & LayerNorm

• Residual: output = LayerNorm(x + sublayer(x)) (Pre-Norm) или output = x + sublayer(LayerNorm(x)) (Post-Norm).
• LayerNorm нормализует по последней оси (d_model), добавляет learnable scale и shift.

После N слоёв (например, 32 для LLaMA-7B) получаем скрытое представление (batch, seq, d_model).

5. Final LayerNorm

Перед проекцией на словарь применяется ещё один слой LayerNorm (иногда его называют final norm). Это стабилизирует выходные представления перед линейным преобразованием.

6. LM Head: проекция на словарь

LM head — линейный слой без bias (или с bias) размерности d_model × vocab_size. Он преобразует скрытое состояние последнего токена (или всех токенов) в логиты (logits) — не нормализованные оценки для каждого токена словаря.

• Вход: (batch, seq, d_model).
• Выход: (batch, seq, vocab_size).
• Для генерации следующего токена берётся логит последней позиции (batch, 1, vocab_size).

Weight tying (если используется): веса LM head равны транспонированной матрице эмбеддингов.

7. Softmax: логиты → вероятности

Softmax преобразует логиты в вероятности:

P(token_i) = exp(logit_i / T) / Σ_j exp(logit_j / T)

• T — температура (temperature). При T=1 — стандартный softmax; T→0 — argmax; T>1 — более равномерное распределение.
• Выход: распределение вероятностей (batch, vocab_size).

Top-k / top-p sampling применяются после softmax для отбора кандидатов.

8. Численный пример (упрощённый)

Пусть:

• vocab_size = 4 (токены: A, B, C, D)
• d_model = 2
• seq_len = 1 (один токен на входе)

1. Токенизация: "B" → id=1.
2. Embedding: [0.1, 0.5].
3. Positional encoding (RoPE): [0.1, 0.5] (без изменений для pos=0).
4. Один слой трансформера: attention + FFN → [0.3, -0.2].
5. Final LayerNorm → [0.6, -0.4].
6. LM head (матрица 4×2): логиты [0.8, 1.2, -0.5, 0.3].
7. Softmax (T=1): вероятности [0.27, 0.40, 0.07, 0.16]. Наиболее вероятен токен B (id=1).

9. Практические аспекты

• Batch inference: forward pass обрабатывает несколько последовательностей одновременно, маскируя паддинг.
• KV cache: при авторегрессивной генерации кэшируются K и V из предыдущих шагов для ускорения.
• Flash Attention: оптимизированная реализация attention, снижающая использование памяти.
• Precision: часто используется float16 или bfloat16 для экономии памяти и ускорения.

Пет-проект для закрепления

Задача Написать минимальный forward pass для однослойного трансформера с RoPE и softmax на PyTorch, используя случайные веса.

Инструменты Python, PyTorch, NumPy.

Шаги:

1. Реализовать токенизацию (простой словарь из 10 токенов).
2. Создать embedding matrix (10×8).
3. Реализовать RoPE (функция rotate_half).
4. Реализовать один слой: Multi-Head Self-Attention (2 головы) + FFN (d_ff=32) с Pre-LN.
5. Добавить final LayerNorm и LM head (8→10).
6. Прогнать последовательность из 3 токенов, получить вероятности для следующего токена.
7. Сравнить с готовой реализацией (например, из transformers).

Ожидаемый результат Рабочий скрипт, который выводит распределение вероятностей и позволяет менять температуру.

Связь с другими вопросами

Навигация

• Предыдущий: 676
• Следующий: 678
• Индекс: 00. Индекс разборов