Механизм внимания и архитектура трансформера

Трансформер (Vaswani et al., «Attention Is All You Need», 2017) — революционная архитектура, устранившая рекуррентность в обработке последовательностей. Self-attention позволяет каждому элементу взаимодействовать с любым другим напрямую, без рекуррентных узких мест. Трансформеры сейчас доминируют в NLP, зрении, белковой биологии и аудио.

Механизм внимания (Attention)

Интуиция: При переводе «The animal didn't cross the street because it was too tired» → нужно понять, что «it» = «animal». RNN «забывает» дальние зависимости. Attention позволяет явно «смотреть» на нужные позиции при генерации каждого токена.

Scaled Dot-Product Attention:

Attention(Q, K, V) = softmax(QKᵀ/√d_k) · V

Q (queries) ∈ ℝ^{n×d_k}: «что ищем». K (keys) ∈ ℝ^{m×d_k}: «что предлагают» все позиции. V (values) ∈ ℝ^{m×d_v}: «что берём» из выбранных позиций.

Шаги: (1) QKᵀ ∈ ℝ^{n×m}: скалярные произведения каждой пары запрос-ключ. Элемент [i,j]: насколько позиция j «отвечает» на запрос позиции i. (2) /√d_k: масштабирование предотвращает слишком большие значения при высоком d_k (иначе softmax насыщается). (3) softmax: нормировка в распределение (строки суммируются в 1). (4) ·V: взвешенная сумма значений.

Сложность: O(n²·d_k) по памяти и времени — квадратично по длине последовательности.

Multi-Head Attention

Один mechanism внимания видит один аспект. Multi-head: h голов, каждая с разными проекциями (W^Q_i, W^K_i, W^V_i):

MultiHead(Q,K,V) = Concat(head₁,...,headₕ) W^O headᵢ = Attention(QW^Q_i, KW^K_i, VW^V_i)

d_k = d_v = d_model/h. Головы специализируются: одна — на синтаксических зависимостях, другая — на семантической близости, третья — на кореференции. W^O ∈ ℝ^{hd_v × d_model} — проекция конкатенации.

Архитектура трансформера

Positional Encoding: Self-attention инвариантен к перестановке позиций — нет встроенного понятия «порядка». PE(pos, 2i) = sin(pos/10000^{2i/d_model}), PE(pos, 2i+1) = cos(...). Добавляем к эмбеддингам токенов. Learnable PE (BERT): позиционные эмбеддинги обучаются.

Encoder слой: Входные токены → Add(x, MHA(x, x, x)) → LayerNorm → Add(·, FFN(·)) → LayerNorm. FFN = max(0, xW₁ + b₁)W₂ + b₂ (d_ff = 4·d_model). Residual connections + LayerNorm = стабильное обучение.

Decoder слой: Masked self-attention (не видит будущих токенов) + Cross-attention (Q из decoder, K,V из encoder) + FFN. Маскирование для авторегрессивной генерации: при генерации токена t не должен «видеть» токены t+1,...

Pre-training/Fine-tuning: BERT: Masked LM + NSP → fine-tune на NLU задачах (QA, NER, entailment). GPT: LM (next token prediction) → few-shot и zero-shot.

Эффективные альтернативы

Flash Attention (Dao et al., 2022): Не материализует матрицу внимания n×n. Вычисляет тайлами, используя online softmax. 3–6× ускорение, 10–20× меньше памяти (HBM). Позволяет 100K+ контекстное окно.

Sparse Attention (Longformer): Локальное окно ±w + несколько глобальных токенов. O(n·w) вместо O(n²). Efficient для длинных документов.

Grouped Query Attention (GQA, Ainslie et al., 2023): Несколько Q-голов делят K,V. Уменьшает KV-кэш при инференсе. LLaMA-2-70B: 8 KV-голов на 64 Q-головы.

RoPE (Rotary Position Embedding): Позиции кодируются через поворот Q и K: Q(pos) = R(pos)·Q, K(pos) = R(pos)·K. Скалярное произведение Q·K зависит только от (pos₁−pos₂) — относительная позиция. Лучше экстраполируется на длинные контексты.

Численный пример

Трансформер для MT (русский → английский): n=20 токенов, d_model=512, h=8, N=6 слоёв. Параметров: 6×(4·512² + 2·512·2048) = 6×(1.05M + 2.1M) ≈ 19M. KV-cache при инференсе (one token at time): 6 слоёв × 2 (K,V) × 8 голов × d_k = 6×2×512 = 6144 floats per token. При 1000 токенов: 6144×1000×4 байта = 24 МБ (небольшая модель).

Attention паттерны (визуализация): голова 1 — местные грамматические зависимости (прилагательное→существительное). Голова 5 — глобальные семантические (subject→predicate).

Задание: Реализуйте упрощённый трансформер с нуля (PyTorch): Encoder-only, 2 слоя, 4 головы, d_model=64. Задача: sentiment classification IMDB (25K reviews). (1) Токенизация BPE (tokenizers library, vocab=5000). (2) Обучите с нуля на 5000 примеров. (3) Сравните с BERT fine-tune на тех же данных (трансформеры = universal language understanding?). (4) Визуализируйте матрицы внимания для 5 примеров — какие слова «смотрят» на «not», «terrible», «excellent»?