Self-attention a multi-head v transformerech

Mechanismy pozornosti v transformerech: self-attention, cross-attention a multi-head

Mechanismus pozornosti umožňuje modelům umělé inteligence dynamicky vážit důležitost prvků vstupní posloupnosti. Pro každý token se vypočítají váhy na základě matic Query, Key a Value, čímž se vytvářejí kontextualizované reprezentace. Self-attention analyzuje vazby uvnitř jedné posloupnosti, cross-attention spojuje enkodér a dekodér, multi-head attention paralelizuje proces pro různé aspekty dat.

Teorie self-attention

Vstupní posloupnost je zadána maticí $X \in \mathbb{R}^{n \times d}$, kde $n$ je počet tokenů, $d$ je dimenze embedingu. Z $X$ se promítnou matice $Q = XW^Q$, $K = XW^K$, $V = XW^V$ s váhami $W^{Q,K,V} \in \mathbb{R}^{d \times d_k}$.

Matice hodnocení pozornosti $A_{ij} = \frac{Q_i K_j^T}{\sqrt{d_k}}$ měří relevanci $j$-tého tokenu pro $i$-tý. Škálování na $\sqrt{d_k}$ stabilizuje gradienty, čímž zabraňuje vanishing gradients v softmax.

Normalizovaná matice $N = \mathrm{softmax}(A)$ po řádcích dává pravděpodobnostní distribuci. Výstup: $\mathrm{Attention}(Q, K, V) = N V$.

Praktický příklad self-attention

Uvažujme větu "Karina jde do obchodu", tokenizovanou jako ["karina", "jde", "do", "obchod"] s $n=4$. Embeddingy:

| Token | Embedding |

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

| karina | [1, 0] |

| jde | [0, 1] |

| do | [1, 1] |

| obchod | [0.5, 1] |

Matice $X_{4\times2} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \\ 0.5 & 1 \end{bmatrix}$. Pro zjednodušení $W^Q = W^K = W^V = I$, tedy $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}$, dělíme $\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 řádcích: $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}$.

Výstup $N \times V = \begin{bmatrix} 0.73 & 0.68 \\ 0.57 & 0.87 \\ 0.68 & 0.81 \\ 0.63 & 0.84 \end{bmatrix}$. Nové embeddingy berou v úvahu kontext: vektor "jde" nyní odráží pohyb osoby, ne abstraktní pojem času.

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: spojení posloupností

V cross-attention se Query berou z dekodéru ($Q = X_{dec}W^Q$), zatímco Key a Value z enkodéru ($K = X_{enc}W^K$, $V = X_{enc}W^V$). Výpočet je identický: $\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V$.

Rozdíly:

| Komponenta | Self-attention | Cross-attention |

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

| Q | X | X_dec |

| K | X | X_enc |

| V | X | X_enc |

To umožňuje dekodéru zaměřit se na relevantní části kontextu enkodéru.

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: paralelní perspektivy

Multi-head attention provádí $h$ nezávislých výpočtů pozornosti s projekcemi $Q^{(i)} = X W_Q^{(i)}$, $K^{(i)} = X W_K^{(i)}$, $V^{(i)} = X W_V^{(i)}$. Každá hlava: $\mathrm{head}_i = \mathrm{Attention}(Q^{(i)}, K^{(i)}, V^{(i)})$.

Výsledky se zřetězí a promítnou: $\mathrm{MultiHead}(X) = \mathrm{Concat}(head_1, \dots, head_h) W^O$. Velikost hlavy $d_k = d / h$, kde $d$ je 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)
        # Dodatečné vrstvy pro úplnost implementace
        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 je důležité

• Self-attention kontextualizuje tokeny: každý výstupní vektor závisí na celé posloupnosti.
• Škálování $\sqrt{d_k}$ je nezbytné: stabilizuje softmax a gradienty.
• Cross-attention pro seq2seq: spojuje enkodér a dekodér v transformerech.
• Multi-head zachycuje aspekty: gramatiku, sémantiku, vzdálené závislosti.
• PyTorch implementace používají Linear a bmm: pro efektivitu v dávkách.

— Editorial Team