Funkcjonalna rekonstrukcja modeli MLA do dekodowania spekulacyjnego

Naukowcy z arXiv opracowali nową metodę konwersji modeli llenguistycznych z tradycyjnej architektury Multi-Head Attention (MHA) i Grouped-Query Attention (GQA) na bardziej efektywną Multi-Head Latent Attention (MLA), która zmniejsza pamięć wymaganą do przechowywania cache'u KV podczas długotrwałego wnioskowania.

Główny problem polega na tym, że bezpośrednia konwersja między tymi architekturami powoduje błędy w funkcjonowaniu mechanizmu attention, co obniża zgodność między proponowanymi tokenami a weryfikacją w dekodowaniu spekulacyjnym - technice przyspieszającej generowanie tekstu. Zamiast tradycyjnego podejścia opartego na kompresji cache'u, naukowcy sformułowali zadanie jako funkcjonalną rekonstrukcję. Ich metoda end-to-end (E2E) optymalizuje każdy moduł attention w skonwertowanym modelu, aby precyzyjnie odtworzyć odpowiedź oryginalnego MHA/GQA na stanach ukrytych z danych kalibracyjnych.

Metoda jest konwersje-agnostyczna i nie wymaga supervisji od modelu weryfikacyjnego ani dostępu do jego logitów. Badacze testowali rozwiązanie na 192 różnych konfiguracjach, obejmujących pary modeli Llama i Qwen, dwie popularne konwertery (TransMLA i MHA2MLA), dwa backendy (HuggingFace i vLLM) oraz cztery zadania na zbiorach 200 promptów. Wyniki pokazują znaczną poprawę w akceptacji tokenów propozycji przy zachowaniu efektywności cache'u nowych modeli MLA.