07 · Architektury · 4 min czytania · Interaktywne · aktualizacja
Czym jest uwaga wielogłowa (multi-head attention) i po co transformerowi wiele głów?
W skrócie
Uwaga wielogłowa uruchamia kilka niezależnych uwag równolegle w mniejszych podprzestrzeniach, więc każdy token może naraz śledzić różne relacje w tekście.
Co to jest
Uwaga wielogłowa (multi-head attention) to wariant mechanizmu uwagi, w którym zamiast jednej dużej operacji uwagi wykonuje się h mniejszych — tzw. głów — równolegle, każdą z własnymi projekcjami zapytań, kluczy i wartości. Wyniki głów są sklejane w jeden wektor i przepuszczane przez wspólną macierz wyjściową. Każda głowa może się nauczyć zwracać uwagę na coś innego.
Intuicja: czytając zdanie, równocześnie śledzisz kilka rzeczy — kto jest podmiotem, do czego odnosi się zaimek, które słowo stoi tuż obok, jaki jest czas gramatyczny. Jedna głowa uwagi musi rozdzielić „uwagę” sumującą się do 1, więc patrząc w dwa miejsca, dostaje ich rozmytą średnią. Kilka głów to kilku czytelników, każdy z własnym pytaniem, którzy na koniec składają notatki w jedną całość.
Uwagę wielogłową wprowadzili Vaswani i współpracownicy w 2017 roku jako część transformera. Jest w każdym współczesnym modelu językowym i wizyjnym opartym na transformerze.
Mechanizm — dlaczego tak działa
Dla wymiaru modelu d i h głów każda głowa dostaje wymiar d_k = d / h. Głowa i liczy Q_i = X·W_Q^i, K_i = X·W_K^i, V_i = X·W_V^i (macierze rozmiaru d×d_k) i zwykłą skalowaną uwagę: head_i = softmax(Q_i·K_iᵀ / √d_k)·V_i. Następnie MultiHead(X) = Concat(head_1, …, head_h)·W_O, gdzie W_O ma rozmiar d×d.
Kluczowa obserwacja: wiele głów nie kosztuje więcej parametrów niż jedna duża. Osiem projekcji 512×64 to tyle samo liczb co jedna 512×512. Zmienia się tylko struktura — zamiast jednego rozkładu wag uwagi na token mamy h niezależnych rozkładów. Koszt obliczeń jest podobny, choć macierzy uwagi n×n do przechowania jest h razy więcej.
Dlaczego to pomaga? Softmax w jednej głowie tworzy jeden rozkład, a wynik to jedna średnia ważona. Jeśli token potrzebuje informacji z dwóch różnych miejsc — np. czasownik potrzebuje i podmiotu, i dopełnienia — jedna głowa musi albo wybrać jedno, albo je uśrednić, a uśrednienie zaciera, co skąd pochodzi. Vaswani i in. piszą wprost, że w pojedynczej głowie uśrednianie to utrudnia. Wiele głów pozwala każdej skupić się ostro na innej relacji, a W_O uczy się te informacje łączyć.
Badania nad wytrenowanymi modelami pokazują, że głowy często się specjalizują: jedne patrzą na poprzedni token, inne łączą zaimki z rzeczownikami, inne śledzą relacje składniowe. Specjalizacja nie jest jednak zaprogramowana ani gwarantowana. Okazało się też, że wiele głów jest zbędnych: Michel, Levy i Neubig (2019) pokazali, że po treningu sporą część głów można usunąć przy niewielkiej stracie jakości. Głowy są więc częściowo nadmiarowe — co prawdopodobnie ułatwia trening.
W dużych modelach przy generowaniu tekstu wąskim gardłem jest pamięć na klucze i wartości poprzednich tokenów (pamięć podręczna KV). Stąd warianty: multi-query attention (wszystkie głowy dzielą jedne klucze i wartości) i grouped-query attention (grupy głów dzielą klucze i wartości), które zmniejszają tę pamięć przy niewielkiej stracie jakości.
Na przykładzie
Zdanie „kot pije mleko”. Klucze (dwuwymiarowe): kot = [4, 0], pije = [0, 0], mleko = [0, 4]; wartości: kot = [1, 0], pije = [0, 0], mleko = [0, 1]. Token „pije” potrzebuje i podmiotu, i dopełnienia. Głowa 1 z zapytaniem [1, 0] daje wagi 0,89 (kot), 0,05 (pije), 0,05 (mleko) i wynik [0,89; 0,05]. Głowa 2 z zapytaniem [0, 1] daje wagi 0,05, 0,05, 0,89 i wynik [0,05; 0,89]. Po sklejeniu „pije” dostaje wektor [0,89; 0,05; 0,05; 0,89] — w jednej połowie informację o podmiocie, w drugiej o dopełnieniu. Jedna głowa z zapytaniem [1, 1] musi rozdzielić uwagę: 0,49, 0,03, 0,49, a wynik [0,49; 0,49] jest rozmytą mieszanką, z której nie wiadomo już, co było podmiotem.
Rzeczywiste konfiguracje: oryginalny transformer (wersja base) ma d = 512 i 8 głów po 64 wymiary; GPT-2 small ma d = 768 i 12 głów po 64 wymiary; BERT-base — tak samo 12 głów po 64; GPT-3 175B ma d = 12 288 i 96 głów po 128 wymiarów. W GPT-2 small projekcje Q, K, V i W_O jednej warstwy to 4·768² = 2 359 296 wag — tyle samo, ile miałaby jedna głowa o wymiarze 768.
W praktyce
- PyTorch:
nn.MultiheadAttention(embed_dim=768, num_heads=12, batch_first=True);embed_dimmusi być podzielne przeznum_heads. - Typowy wymiar głowy to 64 lub 128; liczbę głów dobiera się tak, by d / h trafiało w ten zakres.
- Wagi uwagi poszczególnych głów można odczytać (
need_weights=True, average_attn_weights=False) i zwizualizować, np. biblioteką BertViz. - Przy dużych modelach rozważ grouped-query attention — mniejsza pamięć KV i szybsze generowanie.
- Częsty błąd: przy własnej implementacji pomylenie kolejności
reshapeitransposeprzy dzieleniu na głowy — kod działa, ale głowy mieszają wymiary różnych tokenów.
Najczęstsze pytania
- Czy więcej głów zawsze znaczy lepiej?
- Nie. Przy stałym d więcej głów oznacza mniejszy wymiar każdej z nich, a zbyt małe głowy (np. 8–16 wymiarów) słabiej rozróżniają tokeny. Wartości 8–16 głów przy d rzędu 512–1024 to sprawdzony kompromis.
- Czy głowy same się specjalizują?
- Często tak — znajdowano głowy śledzące poprzedni token, składnię czy odwołania zaimków. Ale nie jest to wymuszone: część głów robi podobne rzeczy, a część da się usunąć bez dużej straty.
- Czy uwaga wielogłowa jest droższa niż zwykła?
- Pod względem parametrów — nie, bo projekcje mają łącznie ten sam rozmiar. Pod względem pamięci na wagi uwagi — tak, bo każda głowa ma własną macierz n×n.
Źródła
- Vaswani i in. „Attention Is All You Need”, NeurIPS 2017, arXiv:1706.03762.
- Michel, Levy, Neubig „Are Sixteen Heads Really Better than One?”, NeurIPS 2019.
- Voita, Talbot, Moiseev, Sennrich, Titov „Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned”, ACL 2019.
- Ainslie i in. „GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”, EMNLP 2023.
- Zhang i in. „Dive into Deep Learning”, d2l.ai, rozdz. 11.5 („Multi-Head Attention”).