Multi-Query Attention

Model łączy aktualny fragment tekstu z wcześniejszą rozmową za pomocą kilku równoległych obliczeń. Takie obliczenia nazywa się głowicami uwagi, czyli attention. Jeśli każda głowica przechowuje własne opisy całej wcześniejszej rozmowy, rośnie koszt pamięci i jej odczytu. Czy można współdzielić część danych, zachowując różne sposoby ich odczytywania?

Multi-Query Attention (MQA) pozostawia wiele query — opisów potrzeb aktualnych pozycji — ale daje im wspólny zestaw key i value. Key służy do dopasowania, a value zawiera informacje do przekazania. Wszystkie są listami liczb.

Dwie głowice mogą używać tego samego zbioru key/value, a mimo to inaczej rozłożyć udziały, bo mają różne query. Wspólne dane nie wymuszają identycznego wyniku.

Oszczędność dotyczy m.in. KV Cache, czyli przechowywanych opisów wcześniejszego tekstu. Zmiana może jednak wpływać na jakość i nie oznacza proporcjonalnego przyspieszenia całego modelu. Grouped-Query Attention stosuje współdzielenie wewnątrz mniejszych grup.

Mechanizm i szczegóły

W zwykłym Multi-Head Attention każda głowa ma własne projekcje query, key i value. MQA usuwa wymiar głów z K/V oraz ich projekcji, zachowując go w query. Tak definiuje wariant Noam Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, §3.

Ta sama pamięć, inne pytanie

Własny przykład ma dwa klucze: $k_1=(1,0)$ i $k_2=(0,1)$. Dwie query heads otrzymują odpowiednio $q_1=(2,0)$ i $q_2=(0,2)$. Scaled Dot-Product Attention dzieli iloczyny przez $\sqrt{2}$ i stosuje Softmax. Pierwsza głowa daje kluczom wagi około 0,804 i 0,196; druga odwraca te proporcje. Przy wspólnych $v_1=(1,0)$ i $v_2=(0,1)$ są to zarazem współrzędne dwóch różnych wyników.

Skąd oszczędność

KV Cache przechowuje wcześniejsze K/V. Przy ośmiu query heads MHA zapisuje osiem zestawów K/V, a MQA jeden. Dla tych samych długości sekwencji, wymiarów głowy i precyzji same tensory cache są wtedy osiem razy mniejsze. To rachunek rozmiaru, nie ośmiokrotne przyspieszenie całego modelu.

Analiza Shazeera w §2.4.1 i §3.1 wiąże tę zmianę z transferem pamięci podczas kolejnych kroków generowania. Zmienia się jednak architektura i jej możliwości reprezentacji; nie jest to bezstratne usunięcie identycznych kopii z dowolnego wytrenowanego MHA. Grouped-Query Attention zachowuje kilka zestawów K/V jako wariant pośredni.

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