ML Atlas

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.

Ta ilustracja działa w przeglądarce z włączonym JavaScriptem: samouwaga w zdaniu „Kot zjadł rybę, bo był głodny” na ręcznie ustawionych wagach: „był” kieruje najwięcej uwagi na „Kot”.

W praktyce

  • PyTorch: nn.MultiheadAttention(embed_dim=768, num_heads=12, batch_first=True); embed_dim musi być podzielne przez num_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 reshape i transpose przy 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”).

Zobacz też