На чем основан механизм внимания в архитектуре Transformer?

Механизм внимания (Attention) в Transformer основан на операции Scaled Dot-Product Attention. Его цель — вычислить контекстное представление для каждого элемента последовательности как взвешенную сумму всех элементов, где веса определяются их взаимной релевантностью.

Формула внимания: Attention(Q, K, V) = softmax( (Q * K^T) / sqrt(d_k) ) * V

Компоненты:

1. Q (Query), K (Key), V (Value): Линейные проекции исходных эмбеддингов. Каждый токен получает свои векторы запроса, ключа и значения.
2. *Q K^T:** Скалярное произведение запроса с каждым ключом измеряет схожесть (совместимость).
3. Масштабирование на sqrt(d_k): Стабилизирует градиенты, предотвращая слишком малые значения градиента после softmax при большой размерности ключей.
4. Softmax: Преобразует оценки совместимости в вероятностное распределение весов (сумма = 1).
5. Умножение на V: Выход — взвешенная сумма значений, где веса определяют, насколько каждому токену следует "уделить внимание" другим токенам.

Практическая реализация (PyTorch):

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q, K, V: тензоры формы [batch_size, seq_len, d_model]
    mask: тензор для маскирования (например, для padding или будущих токенов в декодере)
    """
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))

    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9) # Маскированные позиции получают очень низкий вес

    attn_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attn_weights, V)
    return output, attn_weights

Multi-Head Attention: Для повышения выразительности механизм выполняется параллельно в h разных подпространствах (головах), а их выходы конкатенируются и проецируются обратно.