Grouped-Query Attention

Grouped-Query Attention (GQA) dzieli query heads na grupy. Głowy z jednej grupy korzystają ze wspólnej key head i value head, lecz każda nadal oblicza własne wagi attention. Liczba query heads i liczba KV heads są więc odrębnymi parametrami.

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.