Multi-Query Attention

Multi-Query Attention (MQA) zachowuje wiele query heads, ale współdzieli między nimi jedną key head i jedną value head. Każda query head nadal osobno oblicza swoje wagi attention. Wspólne K i V nie oznaczają wspólnego wyniku.

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.