Self-attention i multi-head w transformerach

Mechanizmy uwagi w transformatorach: self-attention, cross-attention i multi-head

Mechanizm uwagi pozwala modelom sztucznej inteligencji dynamicznie oceniać ważność elementów sekwencji wejściowej. Dla każdego tokenu obliczane są wagi na podstawie macierzy Query, Key i Value, tworząc skontekstualizowane reprezentacje. Self-attention analizuje powiązania wewnątrz jednej sekwencji, cross-attention łączy enkoder i dekoder, multi-head attention równolegle przetwarza różne aspekty danych.

Teoria self-attention

Sekwencja wejściowa jest reprezentowana przez macierz $X \in \mathbb{R}^{n \times d}$, gdzie $n$ to liczba tokenów, $d$ to wymiarowość embeddingu. Z $X$ projektowane są macierze $Q = XW^Q$, $K = XW^K$, $V = XW^V$ z wagami $W^{Q,K,V} \in \mathbb{R}^{d \times d_k}$.

Macierz ocen uwagi $A_{ij} = \frac{Q_i K_j^T}{\sqrt{d_k}}$ mierzy relewantność $j$-tego tokena dla $i$-tego. Skalowanie przez $\sqrt{d_k}$ stabilizuje gradienty, zapobiegając vanishing gradients w softmax.

Znormalizowana macierz $N = \mathrm{softmax}(A)$ po wierszach daje rozkład probabilistyczny. Wyjście: $\mathrm{Attention}(Q, K, V) = N V$.

Praktyczny przykład self-attention

Rozważmy zdanie „Karina idzie do sklepu”, tokenizowane jako ["karina", "idzie", "do", "sklepu"] z $n=4$. Embeddingi:

| Token | Embedding |

|---------|--------------|

| karina | [1, 0] |

| idzie | [0, 1] |

| do | [1, 1] |

| sklep | [0.5, 1] |

Macierz $X_{4\times2} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \\ 0.5 & 1 \end{bmatrix}$. Dla uproszczenia $W^Q = W^K = W^V = I$, więc $Q=K=V=X$.

$QK^T = \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}$, dzielimy przez $\sqrt{2} \approx 1.41$:

$A \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}$.

Po softmax po wierszach: $N \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}$.

Wyjście $N \times V = \begin{bmatrix} 0.73 & 0.68 \\ 0.57 & 0.87 \\ 0.68 & 0.81 \\ 0.63 & 0.84 \end{bmatrix}$. Nowe embeddingi uwzględniają kontekst: wektor „idzie” odzwierciedla teraz ruch osoby, a nie abstrakcyjne pojęcie czasu.

import torch  
import torch.nn as nn  

class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.embed_dim = embed_dim
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)
        self.scale = embed_dim ** 0.5

    def forward(self, x):
        Q = self.q_linear(x)
        K = self.k_linear(x)
        V = self.v_linear(x)
        scores = torch.bmm(Q, K.transpose(1, 2)) / self.scale
        attn = torch.softmax(scores, dim=-1)
        out = torch.bmm(attn, V)
        return out

Cross-attention: połączenie sekwencji

W cross-attention Query pochodzą z dekodera ($Q = X_{dec}W^Q$), a Key i Value z enkodera ($K = X_{enc}W^K$, $V = X_{enc}W^V$). Obliczenie jest identyczne: $\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V$.

Różnice:

| Komponent | Self-attention | Cross-attention |

|-----------|----------------|-----------------|

| Q | X | X_dec |

| K | X | X_enc |

| V | X | X_enc |

Pozwala to dekoderowi skupiać się na relewantnych częściach kontekstu enkodera.

import torch  
import torch.nn as nn  

class CrossAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.embed_dim = embed_dim
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)
        self.scale = embed_dim ** 0.5

    def forward(self, x_dec, x_enc):
        Q = self.q_linear(x_dec)
        K = self.k_linear(x_enc)
        V = self.v_linear(x_enc)
        scores = torch.bmm(Q, K.transpose(1, 2)) / self.scale
        attn = torch.softmax(scores, dim=-1)
        out = torch.bmm(attn, V)
        return out

Multi-head attention: równoległe perspektywy

Multi-head attention wykonuje $h$ niezależnych obliczeń uwagi z projekcjami $Q^{(i)} = X W_Q^{(i)}$, $K^{(i)} = X W_K^{(i)}$, $V^{(i)} = X W_V^{(i)}$. Każda głowa: $\mathrm{head}_i = \mathrm{Attention}(Q^{(i)}, K^{(i)}, V^{(i)})$.

Wyniki są konkatenowane i projektowane: $\mathrm{MultiHead}(X) = \mathrm{Concat}(head_1, \dots, head_h) W^O$. Rozmiar głowy $d_k = d / h$, gdzie $d$ to embed_dim.

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)
        # Dodatkowe warstwy dla kompletności implementacji
        self.out = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        B, T, D = x.shape
        Q = self.q(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.k(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.v(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, V).transpose(1, 2).contiguous().view(B, T, D)
        return self.out(out)

Co jest ważne

• Self-attention kontekstualizuje tokeny: każdy wektor wyjściowy zależy od całej sekwencji.
• Skalowanie $\sqrt{d_k}$ jest obowiązkowe: stabilizuje softmax i gradienty.
• Cross-attention dla seq2seq: łączy enkoder i dekoder w transformatorach.
• Multi-head przechwytuje aspekty: gramatyka, semantyka, odległe zależności.
• Implementacje w PyTorch używają Linear i bmm: dla wydajności batchewej.

— Editorial Team