Краткий справочник про внимания (self-attention, cross-attention, multi-head attention) / Хабр

Механизм внимания (Attention) - это метод в искусственном интеллекте, который позволяет нейросети динамически определять, какие части входных данных наиболее важны для текущей задачи. Он работает через вычисление весов важности для разных элементов входа: более важные элементы получают больший вес, а менее важные - меньший. Затем модель формирует взвешенную сумму представлений, создавая новый контекстный вектор.

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

Как работает self-attention (Теория)

Пусть входная последовательность представлена матрицей:

$X \in\mathbb{R}^{n \times d}$

n - число токенов, d - размерность эмбеддинга

Из X получаем три матрицы Query (Q), Key (K) и Value (V), где $Q = XW^Q, \quad K = XW^K, \quad V = XW^V$ где $W^Q, W^K,W^V\in\mathbb{R}^{d \times d_k}$

Attention scores (оценка важности)

Считаем "похожесть" между токенами: $A_{ij} = \frac{Q_iK_j^T}{\sqrt{d_k}}$ где $A \in \mathbb{R}^{n \times n}$ , элемент $A_{ij}$ показывает, насколько i-й токен "смотрит" на j-й.

Делим на $\sqrt{d_k}$ , потому что компоненты Q и K имеют дисперсию 1, то их скалярное произведение имеет дисперсию . При большом softmax насыщается и градиенты исчезают. Деление на $\sqrt{d_k}$ возвращает дисперсию к 1. (допущение)

Нормализуем (то есть применяем softmax каждой строке матрицы A), $N = softmax(A_{ij})$ , каждая строка - распределение вероятностей. Тогда получаем:

$\text{Attention}(Q, K, V) = NV$

То есть,

$\text{Attention}(Q, K, V) = \mathrm{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V$

Как работает self-attention (Пример)

Допустим есть предложение “Карина идет в магазин”. Мы не можем работать с буквами напрямую, поэтому представим текст в виде токенов. Для упрощения будем считать, что 1 токен = 1 слово.

Эмбеддинги

Токенизация. Разбиваем фразу на части: ["карина", "идет", "в", "магазин"]. Здесь (число токенов).

Эмбеддинг. Каждый токен заменяется вектором размерности . Для упрощения, допустим что:

слово	вектор
Карина	[1, 0]
идет	[0, 1]
в	[1, 1]
магазин	[0.5, 1]

Position Encoding (позиционное кодирование). К векторам добавляется информация о позиции, чтобы модель знала, что «карина» стоит на 1-м месте, а «магазин» на 4-м.

Преобразования

Мы получили:

$X_{4 \times 2} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \\ 0.5 & 1 \end{bmatrix}$

Чтобы не утонуть в матрицах, возьмём: , тогда . (однако надо понимать, что матрицы это обучаемые параметры слоя).

Оценка важности

$A = \frac{QK^T}{\sqrt{d_k}}$

Мы имеем (т.к размерность векторов эмбеддинга равна 2). Тогда $\sqrt{d} \approx 1.41$

$QK^T = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \\ 0.5 & 1 \end{bmatrix} \cdot \begin{bmatrix} 1 & 0 & 1 & 0.5 \\ 0 & 1 & 1 & 1 \end{bmatrix} = \begin{bmatrix} 1 & 0 & 1 & 0.5 \\ 0 & 1 & 1 & 1 \\ 1 & 1 & 2 & 1.5 \\ 0.5 & 1 & 1.5 & 1.25 \end{bmatrix}$

Делим на $\sqrt{d} \approx 1.41$

$A = \frac{1}{1.41} \begin{bmatrix} 1 & 0 & 1 & 0.5 \\ 0 & 1 & 1 & 1 \\ 1 & 1 & 2 & 1.5 \\ 0.5 & 1 & 1.5 & 1.25 \end{bmatrix} \approx \begin{bmatrix} 0.71 & 0 & 0.71 & 0.35 \\ 0 & 0.71 & 0.71 & 0.71 \\ 0.71 & 0.71 & 1.42 & 1.06 \\ 0.35 & 0.71 & 1.06 & 0.89 \end{bmatrix}$

Дальше применяем softmax, где $\text{Softmax}(x_i) = \frac{e^{x_i}}{\sum_{j=1}^{n} e^{x_j}}$ . Важно понимать, что softmax применяется к каждой строке отдельно. Исходная 1 строка, где $x_1 = [0.71 \quad 0 \quad 0.71 \quad 0.35]$ тогда подставив в $\text{Softmax}(x_i)$ получим $[0.31 \quad 0.15 \quad 0.31 \quad 0.22]$ . Что это означает для «Карины»? Теперь модель знает, что при формировании смысла слова «Карина» в этом предложении: 31% информации она берет из самой себя, 15% - из слова «идет» и так далее. Подставив каждую строку, получим матрицу

$N = \text{Softmax}(A) \approx \begin{bmatrix} 0.31 & 0.15 & 0.31 & 0.22 \\ 0.14 & 0.29 & 0.29 & 0.29 \\ 0.18 & 0.18 & 0.37 & 0.26 \\ 0.16 & 0.23 & 0.33 & 0.28 \end{bmatrix}$

Последний этап это умножение матрицы весов (результат softmax) на саму информацию, которую несут слова, то есть на матрицу Value (V).

$\text{output} = N \times V$ $\text{Output} = \begin{bmatrix} 0.31 & 0.15 & 0.31 & 0.22 \\ 0.14 & 0.29 & 0.29 & 0.29 \\ 0.18 & 0.18 & 0.37 & 0.26 \\ 0.16 & 0.23 & 0.33 & 0.28 \end{bmatrix} \cdot \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \\ 0.5 & 1 \end{bmatrix} = \begin{bmatrix} 0.73 & 0.68 \\ 0.57 & 0.87 \\ 0.68 & 0.81 \\ 0.63 & 0.84 \end{bmatrix}$

Матрица ‘Output’ это новые эмбеддинги слов после self-attention, то есть

токен	было	стало
Карина	[1, 0]	[0.73, 0.68]
идет	[0, 1]	[0.57, 0.87]
в	[1, 1]	[0.68, 0.81]
магазин	[0.5, 1]	[0.63, 0.84]

Каждый вектор после self-attention уже не описывает слово само по себе, а описывает его с учётом всех остальных слов в предложении.Например, слово «идет» в контексте «Карина идет в магазин» получает представление, связанное с физическим перемещением человека.А в предложении «Время идет быстро» это же слово будет иметь другое контекстное представление, связанное с течением времени и изменением состояния. То есть модель не хранит разные “значения” слова отдельно, а каждый раз формирует новое представление слова на основе контекста.

После self-attention представление каждого токена становится контекстуализированным: вектор теперь зависит от всего предложения и отражает смысл слова с учётом окружающего контекста.

Self-attention на Pytorch

import torch  
import torch.nn as nn  

class SelfAttention(nn.Module):  # класс self-attention слоя
    def __init__(self, embed_dim):  # embed_dim = размер эмбеддинга
        super().__init__()  # инициализация nn.Module
        self.embed_dim = embed_dim  # сохраняем размерность

        self.q_linear = nn.Linear(embed_dim, embed_dim)  # слой для Query
        self.k_linear = nn.Linear(embed_dim, embed_dim)  # слой для Key
        self.v_linear = nn.Linear(embed_dim, embed_dim)  # слой для Value

        self.scale = embed_dim ** 0.5  # коэффициент масштабирования (sqrt(d))

    def forward(self, x):
        """
        x: (batch_size, seq_len, embed_dim)
        """
        Q = self.q_linear(x)  # (B, T, D) — Query
        K = self.k_linear(x)  # (B, T, D) — Key
        V = self.v_linear(x)  # (B, T, D) — Value

        # attention scores
        scores = torch.bmm(Q, K.transpose(1, 2))  # (B, T, T) — QK^T
        scores = scores / self.scale  # нормализация (stability)

        attn = torch.softmax(scores, dim=-1)  # вероятности внимания

        out = torch.bmm(attn, V)  # взвешенная сумма V

        return out # выход

Cross-Attention, Multi-Head Attention

Cross-Attention

Cross-attention - это механизм внимания, где запросы (Q) берутся из одной последовательности, а ключи (K) и значения (V) - из другой. Cross-attention считается так же, как self-attention, но Q берется из другой последовательности.

$Q = X_{dec}W^Q, \quad K = X_{enc}W^K, \quad V = X_{enc}W^V$ , далее

$\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V$

Что меняется по сравнению с self-attention:

компонент	self-attention	cross-attention
Q	X	X_dec
K	X	X_enc
V	X	X_enc

Cross-Attention в Pytoch

import torch  
import torch.nn as nn  

class CrossAttention(nn.Module):  # класс cross-attention слоя
    def __init__(self, embed_dim):  # embed_dim = размер эмбеддинга
        super().__init__()  # инициализация nn.Module
        self.embed_dim = embed_dim  # сохраняем размерность

        self.q_linear = nn.Linear(embed_dim, embed_dim)  # слой для Query (из декодера)
        self.k_linear = nn.Linear(embed_dim, embed_dim)  # слой для Key (из энкодера)
        self.v_linear = nn.Linear(embed_dim, embed_dim)  # слой для Value (из энкодера)

        self.scale = embed_dim ** 0.5  # коэффициент масштабирования (sqrt(d))

    def forward(self, x_dec, x_enc):
        """
        x_dec: (batch_size, seq_len_q, embed_dim) — запрос
        x_enc: (batch_size, seq_len_kv, embed_dim) — контекст
        """
        Q = self.q_linear(x_dec)  # (B, T_q, D) — Query из декодера
        K = self.k_linear(x_enc)  # (B, T_kv, D) — Key из энкодера
        V = self.v_linear(x_enc)  # (B, T_kv, D) — Value из энкодера

        # attention scores
        # (B, T_q, D) x (B, D, T_kv) -> (B, T_q, T_kv)
        scores = torch.bmm(Q, K.transpose(1, 2))  # QK^T — сопоставляем запрос с контекстом
        scores = scores / self.scale  # нормализация 

        attn = torch.softmax(scores, dim=-1)  # вероятности внимания 

        out = torch.bmm(attn, V)  # взвешенная сумма V из энкодера

        return out # выход

Multi-Head Attention

Multi-Head Attention (многоголовое внимание) — это механизм, который запускает несколько attention «параллельно», чтобы модель могла смотреть на последовательность под разными углами. Каждая голова учится фокусироваться на разных аспектах входной последовательности, таких как грамматическая структура, семантическое значение или отношения на расстоянии. Идея: вместо одного набора мы делаем разных наборов (голов), у каждой свои матрицы весов:

$Q^{(i)} = XW_Q^{(i)}, \quad K^{(i)} = XW_K^{(i)}, \quad V^{(i)} = XW_V^{(i)}$

где

Дальше каждая голова считает масштабированное скалярное произведение для внимания (scaled dot-product attention):

$\text{head}_i = \text{Attention}(Q^{(i)}, K^{(i)}, V^{(i)}) = \text{softmax}\left(\frac{Q^{(i)}{K^{(i)}}^T}{\sqrt{d_k}}\right)V^{(i)}$

После этого мы просто «склеиваем» результаты всех голов по последнему измерению и снова линейно преобразуем:

$\text{MultiHead}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O$

где - ещё одна обучаемая матрица.

Multi-Head Attention на PyTorch

import torch
import torch.nn as nn

class MultiHeadSelfAttention(nn.Module):  
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        assert embed_dim % num_heads == 0  # проверка делимости

        self.embed_dim = embed_dim  # размер эмбеддинга
        self.num_heads = num_heads  # число голов
        self.head_dim = embed_dim // num_heads  # размер головы

        self.q = nn.Linear(embed_dim, embed_dim)  
        self.k = nn.Linear(embed_dim, embed_dim)  
        self.v = nn.Linear(embed_dim, embed_dim)  
        self.out = nn.Linear(embed_dim, embed_dim)  # выходная проекция

        self.scale = self.head_dim ** 0.5  # sqrt(d_h)

    def forward(self, x):
        
        B, T, D = x.shape  #batch, seq_len, dim

        Q = self.q(x)  
        K = self.k(x)  
        V = self.v(x)  

        # разбиваем на головы
        Q = Q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)  # (B,h,T,d_h)
        K = K.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)  # (B,h,T,d_h)
        V = V.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)  # (B,h,T,d_h)

        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale  # QK^T / sqrt(d)
        attn = torch.softmax(scores, dim=-1)  # веса внимания

        head_out = torch.matmul(attn, V)  # взвешенная сумма V

        # склеиваем головы обратно 
        head_out = head_out.transpose(1, 2).contiguous().view(B, T, D) 

        return self.out(head_out)  #выход

Примечания:

https://www.mql5.com/ru/articles/8909, https://arxiv.org/abs/1706.03762, https://en.wikipedia.org/wiki/Attention_Is_All_You_Need, https://neerc.ifmo.ru/wiki/index.php?title=Механизм_внимания, https://jalammar.github.io/illustrated-transformer/

Мои контакты: tg: @salyamq2