Nowy Mechanizm Atencji Redukujący Pamięć KV Cache

Wprowadzenie

Wraz z rosnącym wykorzystaniem dużych modeli językowych (LLM) i pojawianiem się nowych paradygmatów wnioskowania, na pierwszy plan wysuwa się wyzwanie efektywnego wnioskowania na dużą skalę. Istotnym wąskim gardłem jest pamięć podręczna Key-Value (KV) w tradycyjnych mechanizmach uwagi, która rozszerza się liniowo wraz z rozmiarem partii i długością sekwencji, stając się "pożeraczem pamięci", który utrudnia skalowanie i rozszerzanie LLM.

Warianty takie jak MQA, GQA i MLA pojawiły się, aby rozwiązać ten problem, ale często mają trudności z utrzymaniem wydajności przy ścisłych ograniczeniach pamięci lub wprowadzają komplikacje, które stanowią wyzwania inżynieryjne i problemy z kompatybilnością.

Multi-matrix Factorization Attention (MFA)

Niedawny artykuł naukowców z Stepes, Uniwersytetu Tsinghua i innych instytucji wprowadza nową architekturę mechanizmu uwagi: Multi-matrix Factorization Attention (MFA) i jego wariant MFA-Key-Reuse (MFA-KR). Ten mechanizm znacząco obniża koszt wnioskowania modeli językowych, jednocześnie poprawiając wydajność.

MFA i MFA-KR nie tylko przewyższają MLA pod względem wydajności, ale także dorównują tradycyjnej wydajności MHA, redukując zużycie pamięci KV Cache nawet o 93,7%.

MFA został zaprojektowany z myślą o prostocie, łatwej reprodukcji, niskiej wrażliwości na hiperparametry i kompatybilności z różnymi metodami Pos-embedding.

Podejście i analiza MFA

Zespół badawczy przeanalizował ogólny projekt i możliwości mechanizmów uwagi, identyfikując dwa krytyczne wymiary związane z pojemnością. Analiza ta doprowadziła do opracowania nowych metod analitycznych i zasad projektowania.

Wprowadzili koncepcję Generalized Multi-Head Attention (GMHA) jako ujednolicające ramy do zrozumienia różnych wariantów MHA.

Zespół zbadał również obliczenia i przechowywanie klucz-wartości z perspektywy wnioskowania i przeanalizował pojemność modelu z perspektywy dekompozycji.

Fully Parameterized Bilinear Attention (FPBA) został ustalony jako teoretyczna górna granica wydajności. Odkryli, że MHA i jego warianty są dekompozycjami niskiego rzędu FPBA.

Porównanie z MQA i MLA

Analiza skupiła się na dwóch reprezentatywnych schematach ulepszeń: Multi-Query Attention (MQA) i Multi-Head Latent Attention (MLA).

MQA stosuje bardziej agresywną strategię współdzielenia parametrów, gdzie wszystkie głowy uwagi współdzielą ten sam zestaw parametrów klucz-wartość. Zmniejsza to zużycie pamięci, ale może wpływać na ekspresyjność modelu.
MLA wprowadza wspólną przestrzeń utajoną do kompresji parametrów, ale rzeczywista moc ekspresyjna jest ograniczona najmniejszym wymiarem, co oznacza, że zwiększanie wymiarów pośrednich nie poprawia znacząco wydajności.

Kluczowe Innowacje MFA

Rozwój MFA był podyktowany celem stworzenia mechanizmu uwagi, który minimalizuje zużycie zasobów, zbliżając się jednocześnie do teoretycznych granic wydajności. Projekt MFA obejmuje trzy kluczowe innowacje:

Znaczące zwiększenie liczby i wymiaru głów uwagi w celu zmaksymalizowania pojemności modelu.
Zastosowanie agresywnej strategii rozkładu niskiego rzędu w celu utrzymania efektywności parametrów przy jednoczesnym zwiększeniu liczby i wymiarów głów uwagi.
Wykorzystanie pojedynczej głowy klucz-wartość w celu zminimalizowania zużycia pamięci, nawet przy zwiększonej złożoności modelu.

Pomiar i porównanie pojemności

Aby dokładniej przeanalizować MFA i inne mechanizmy uwagi, zespół wprowadził dwie kluczowe metryki:

Total Effective Rank (TER): Iloczyn liczby głów uwagi i rzędu faktoryzacji na głowę (FRH).
Shared Latent Subspace Dimension (SLSD): Wymiar ukrytej przestrzeni współdzielonej przez wszystkie głowy uwagi.

MFA osiąga wyższe SLSD i TER w porównaniu z MQA.

W porównaniu z MLA, MFA osiąga mniejszy rozmiar pamięci podręcznej KV i wyższy TER przy podobnych budżetach parametrów, zachowując porównywalny SLSD.

W porównaniu z tradycyjnym MHA, MFA ma wyższy TER, mimo że jego SLSD jest mniejszy.

Wyniki eksperymentalne

Przeprowadzono szeroko zakrojone eksperymenty w celu oceny wydajności nowej architektury na większych skalach, testując modele o parametrach od 1B do 7B i dane treningowe od 10B do 1T.

MFA wykazało zdolności skalowania porównywalne z tradycyjnym MHA, utrzymując doskonałą wydajność nawet na większych skalach.

Chociaż MFA-KR wykazał nieco niższą wydajność, jego trend skalowania był zgodny z MHA. Zalety MFA i MFA-KR w zakresie oszczędzania pamięci rosły wraz z rozmiarem modelu, przy czym MFA osiągnął 87,5% oszczędności pamięci, a MFA-KR zmniejszył zużycie pamięci do 6,25% w największej skali.

Badania ablacjne

Badania ablacjne potwierdziły skuteczność MFA i MFA-KR. Ich zalety wydajnościowe zostały również potwierdzone w różnych głównych metodach kodowania pozycyjnego.

Perspektywy

MFA oferuje znaczące ulepszenia dzięki prostemu projektowi, skutecznie rozwiązując problem wąskiego gardła pamięci we wnioskowaniu LLM bez dodawania dodatkowej złożoności inżynieryjnej. Integruje się bezproblemowo z istniejącym ekosystemem Transformer, przyspieszając zastosowanie LLM w różnych scenariuszach.