Grouped-Query Attention

Model odczytuje wcześniejszy tekst na kilka sposobów równolegle. Każde takie obliczenie jest głowicą uwagi, czyli attention. Chcesz zmniejszyć pamięć zajmowaną przez te głowice, ale pełne współdzielenie jednego zestawu danych może zbyt mocno ograniczać model. Możliwy jest wariant pośredni: część głowic korzysta ze wspólnego zestawu, a inne z innego.

Grouped-Query Attention (GQA) dzieli głowice query na grupy współdzielące key i value. Query to liczbowy opis służący do odczytu kontekstu, key do dopasowania, a value do przekazania informacji. Każda głowica nadal oblicza własne udziały.

Osiem głowic query podzielonych po cztery na dwie grupy oznacza dwa zestawy key/value, a nie osiem. Jedna grupa daje MQA; po jednej grupie na głowicę odpowiada zwykłemu MHA.

Mniej zestawów zmniejsza koszt przechowywania KV przy zgodnych wymiarach. To kompromis konstrukcji i jakości, a nie przełącznik gwarantujący określoną szybkość. Czas zależy również od sprzętu i pozostałych obliczeń.

Mechanizm i szczegóły

Dla $H$ query heads i $G$ grup otrzymujemy $G$ zestawów K/V. Gdy $G=1$, jest to Multi-Query Attention. Gdy $G=H$, wracamy do Multi-Head Attention. Tę rodzinę i oznaczenie GQA-g opisują Ainslie et al., GQA, §2.2.

Policz pamięć, zanim zgadniesz szybkość

Własny rachunek dla KV Cache zakłada batch size 1, 32 warstwy, 1024 zapisane tokeny, wymiar każdej key/value head równy 128 i 2 bajty na liczbę. Same elementy K/V zajmują:

$$2\cdot32\cdot1024\cdot G\cdot128\cdot2\ \text{bajtów}$$

Pierwsza dwójka oznacza K i V. Dla ośmiu query heads wynik wynosi 128 MiB w MHA, 32 MiB dla dwóch KV heads i 16 MiB w MQA. MiB to $2^{20}$ bajtów. Nie liczymy wag modelu, buforów, wyrównania pamięci ani replikacji między urządzeniami.

Redukcja liczby KV heads nie skraca tu kontekstu i nie usuwa query heads. Ogranicza liczbę osobnych reprezentacji K/V. Rachunek pamięci nie przewiduje jakości ani czasu wykonania; zależą one również od treningu, sprzętu i implementacji.

Przy konwersji checkpointu autorzy uśredniali projekcje K/V wewnątrz grup, a następnie kontynuowali Pre-training (§2.1–2.2). Wyniki §3 dotyczą konkretnych wariantów T5 i zadań. Nie ustanawiają uniwersalnej liczby grup ani gwarancji zachowania jakości przy samym przełączeniu konfiguracji.

Wykorzystuję treści generowane przez AI jako część mojego codziennego procesu nauki.