Wzrost popularności Mixture-of-Experts dla wydajnych dużych modeli językowych

W świecie przetwarzania języka naturalnego (NLP) pogoń za budowaniem coraz większych i bardziej zaawansowanych modeli językowych była napędzana przez wiele ostatnich osiągnięć. Jednakże, gdy te modele rosną w rozmiarze, wymagania obliczeniowe dla szkolenia i inferencji stają się coraz bardziej wymagające, napotykając na ograniczenia dostępnych zasobów sprzętowych.

Wkraczają Mixture-of-Experts (MoE), technika, która obiecuje złagodzić ten ciężar obliczeniowy, umożliwiając szkolenie większych i bardziej potężnych modeli językowych. Poniżej omówimy MoE, jego pochodzenie, mechanizm działania i zastosowania w modelach językowych opartych na transformatorach.

Pochodzenie Mixture-of-Experts

Pojęcie Mixture-of-Experts (MoE) można wywieść z wczesnych lat 90-tych, kiedy badacze zbadali ideę obliczeń warunkowych, gdzie części sieci neuronowej są selektywnie aktywowane na podstawie danych wejściowych. Jedną z pionierskich prac w tej dziedzinie była praca “Adaptive Mixture of Local Experts” autorstwa Jacobsa i in. z 1991 roku, która zaproponowała nadzorowaną ramę szkoleniową dla zespołu sieci neuronowych, z których każda specjalizowała się w innym obszarze przestrzeni wejściowej.

Podstawowa idea za MoE polega na tym, że mamy wiele “ekspertów” sieci, z których każdy jest odpowiedzialny za przetwarzanie podzbioru danych wejściowych. Mechanizm bramkowy, zwykle sieć neuronowa, określa, który ekspert (lub eksperci) powinien przetworzyć dane wejściowe. Ten podejście pozwala modelowi alokować swoje zasoby obliczeniowe bardziej efektywnie, aktywując tylko odpowiednich ekspertów dla każdego wejścia, zamiast zatrudniania pełnej pojemności modelu dla każdego wejścia.

Przez lata różni badacze rozwijali i rozszerzali ideę obliczeń warunkowych, co doprowadziło do rozwoju takich rzeczy, jak hierarchiczne MoE, przybliżenia niskiego rzędu dla obliczeń warunkowych i technik do oszacowania gradientów za pomocą neuronów stochastycznych i funkcji aktywacyjnych z progowaniem.

Mixture-of-Experts w transformatorach

Mixture of Experts

Chociaż idea MoE była znana od dekad, jej zastosowanie w modelach językowych opartych na transformatorach jest stosunkowo nowe. Transformatorzy, które stały się standardem dla modeli językowych, składają się z wielu warstw, z których każda zawiera mechanizm uwagi i gęstą sieć neuronową (FFN).

Kluczowa innowacja w zastosowaniu MoE w transformatorach polega na zastąpieniu gęstych warstw FFN rzadkimi warstwami MoE, z których każda składa się z wielu ekspertów FFN i mechanizmu bramkowego. Mechanizm bramkowy określa, który ekspert (lub eksperci) powinien przetworzyć dane wejściowe, umożliwiając modelowi selektywnie aktywować tylko podzbiór ekspertów dla danego ciągu wejściowego.

Jedną z pierwszych prac, która zademonstrowała potencjał MoE w transformatorach, była praca “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer” autorstwa Shazeera i in. z 2017 roku. Praca ta wprowadziła pojęcie warstwy MoE z bramkowaniem, która zastosowała mechanizm bramkowy, który dodał rozproszenie i szum do procesu wyboru ekspertów, zapewniając, że tylko podzbiór ekspertów był aktywowany dla każdego wejścia.

Od tego czasu wiele innych prac rozwinęło zastosowanie MoE w transformatorach, rozwiązując wyzwania takie, jak niestabilność szkolenia, równoważenie obciążenia i efektywna inferencja. Godne uwagi przykłady obejmują Switch Transformer (Fedus i in., 2021), ST-MoE (Zoph i in., 2022) i GLaM (Du i in., 2022).

Korzyści Mixture-of-Experts dla modeli językowych

Główną korzyścią zastosowania MoE w modelach językowych jest możliwość zwiększenia rozmiaru modelu przy zachowaniu relatywnie stałych kosztów obliczeniowych podczas inferencji. Dzięki selektywnemu aktywowaniu tylko podzbioru ekspertów dla każdego tokena wejściowego, modele MoE mogą osiągnąć wyrażalność dużych gęstych modeli przy znacznie mniejszych wymaganiach obliczeniowych.

Na przykład, rozważmy model językowy z gęstą warstwą FFN o 7 miliardach parametrów. Jeśli zastąpimy tę warstwę warstwą MoE składającą się z ośmiu ekspertów, z których każdy ma 7 miliardów parametrów, łączna liczba parametrów wzrośnie do 56 miliardów. Jednak podczas inferencji, jeśli aktywujemy tylko dwóch ekspertów na token, koszt obliczeniowy jest równy kosztowi modelu gęstego o 14 miliardach parametrów, ponieważ wykonujemy dwie mnożenia macierzy o 7 miliardach parametrów.

Ta efektywność obliczeniowa podczas inferencji jest szczególnie cenna w scenariuszach wdrożeniowych, gdzie zasoby są ograniczone, takie jak urządzenia mobilne lub środowiska obliczeń brzegowych. Dodatkowo, zmniejszone wymagania obliczeniowe podczas szkolenia mogą prowadzić do znacznych oszczędności energii i mniejszego śladu węglowego, co jest zgodne z rosnącym naciskiem na zrównoważone praktyki AI.

Wyzwania i rozważania

Chociaż modele MoE oferują kuszące korzyści, ich przyjęcie i wdrożenie wiążą się z kilkoma wyzwaniami i rozważaniami:

1. Niestabilność szkolenia: Modele MoE są znane jako bardziej podatne na niestabilność szkolenia w porównaniu z ich gęstymi odpowiednikami. Problem ten wynika z rzadkiego i warunkowego charakteru aktywacji ekspertów, co może prowadzić do trudności w propagacji gradientów i zbieżności. Techniki takie, jak router z-loss (Zoph i in., 2022), zostały zaproponowane, aby złagodzić te niestabilności, ale nadal potrzebne są dalsze badania.
2. Dostosowanie i przeuczenie: Modele MoE mają tendencję do przeuczenia się podczas dostosowania, szczególnie gdy zadanie downstream ma stosunkowo mały zestaw danych. To zachowanie jest przypisywane zwiększonej pojemności i rzadkości modeli MoE, co może prowadzić do nadmiernego specjalizowania się w danych szkoleniowych. Ostrożne regularizacja i strategie dostosowania są wymagane, aby złagodzić ten problem.
3. Wymagania pamięciowe: Chociaż modele MoE mogą zmniejszyć koszty obliczeniowe podczas inferencji, często mają wyższe wymagania pamięciowe w porównaniu z gęstymi modelami podobnego rozmiaru. Wynika to z faktu, że wszystkie wagi ekspertów muszą być załadowane do pamięci, nawet jeśli tylko podzbiór jest aktywowany dla każdego wejścia. Ograniczenia pamięciowe mogą ograniczyć skalowalność modeli MoE na urządzeniach z ograniczonymi zasobami.
4. Równoważenie obciążenia: Aby osiągnąć optymalną efektywność obliczeniową, konieczne jest zrównoważenie obciążenia między ekspertami, zapewniając, że żaden ekspert nie jest przeciążony, podczas gdy inni pozostają niedostatecznie wykorzystani. To zrównoważenie jest zwykle osiągane za pomocą dodatkowych strat podczas szkolenia i starannego dostosowania czynnika pojemności, który określa maksymalną liczbę tokenów, które mogą być przypisane do każdego eksperta.
5. Nadmiar komunikacyjny: W scenariuszach szkolenia i inferencji rozproszonych, modele MoE mogą wprowadzić dodatkowy nadmiar komunikacyjny z powodu potrzeby wymiany informacji o aktywacji i gradientach między ekspertami znajdującymi się na różnych urządzeniach lub przyspiesznikach. Efektywne strategie komunikacji i świadoma projektowanie modelu są niezbędne, aby złagodzić ten nadmiar.

Pomimo tych wyzwań, potencjalne korzyści modeli MoE w umożliwieniu większych i bardziej zaawansowanych modeli językowych spowodowały znaczne wysiłki badawcze, aby rozwiązać i złagodzić te problemy.

Przykład: Mixtral 8x7B i GLaM

Aby zilustrować praktyczne zastosowanie MoE w modelach językowych, rozważmy dwa godne uwagi przykłady: Mixtral 8x7B i GLaM.

Mixtral 8x7B jest wariantem MoE modelu językowego Mistral, opracowanym przez Anthropic. Składa się z ośmiu ekspertów, z których każdy ma 7 miliardów parametrów, co daje łącznie 56 miliardów parametrów. Jednak podczas inferencji aktywowane są tylko dwa eksperci na token, efektywnie zmniejszając koszt obliczeniowy do poziomu modelu gęstego o 14 miliardach parametrów.

Mixtral 8x7B wykazał imponujące wyniki, przewyższając model Llama o 70 miliardach parametrów, oferując przy tym znacznie szybsze czasy inferencji. Wersja dostosowana do instrukcji Mixtral 8x7B, o nazwie Mixtral-8x7B-Instruct-v0.1, została również wydana, jeszcze bardziej zwiększając jego możliwości w suivaniu naturalnych instrukcji językowych.

Innym godnym uwagi przykładem jest GLaM (Google (GOOGL ) Language Model), duży model MoE opracowany przez Google. GLaM wykorzystuje architekturę transformatora tylko-dekodera i został przeszkolony na ogromnym zestawie danych o 1,6 bilionie tokenów. Model osiąga imponujące wyniki w ocenach few-shot i one-shot, dorównując jakości GPT-3, przy użyciu tylko jednej trzeciej energii wymaganej do szkolenia GPT-3.

Sukces GLaM można przypisać jego efektywnej architekturze MoE, która umożliwiła szkolenie modelu o ogromnej liczbie parametrów przy zachowaniu rozsądnych wymagań obliczeniowych. Model ten również zademonstrował potencjał modeli MoE, aby być bardziej efektywnymi energetycznie i środowiskowo zrównoważonymi w porównaniu z ich gęstymi odpowiednikami.

Architektura Grok-1

GROK MIXTURE OF EXPERT

Grok-1 to model transformatora oparty na MoE o unikalnej architekturze zaprojektowanej w celu maksymalizacji efektywności i wydajności. Zajrzyjmy do kluczowych specyfikacji:

1. Parametry: Z imponującymi 314 miliardami parametrów, Grok-1 jest największym otwartym modelem LLM do tej pory. Jednak dzięki architekturze MoE, tylko 25% wag (około 86 miliardów parametrów) są aktywne w dowolnym momencie, zwiększając możliwości przetwarzania.
2. Architektura: Grok-1 wykorzystuje architekturę Mixture-of-8-Experts, z każdym tokenem przetwarzanym przez dwóch ekspertów podczas inferencji.
3. Warstwy: Model składa się z 64 warstw transformatora, z których każda zawiera mechanizm uwagi i gęstą sieć neuronową.
4. Tokenizacja: Grok-1 wykorzystuje tokenizator SentencePiece o rozmiarze słownika 131 072 tokenów.
5. Wstawienia i kodowanie pozycyjne: Model posiada 6 144-wymiarowe wstawienia i wykorzystuje kodowanie pozycyjne obrotowe, umożliwiając bardziej dynamiczną interpretację danych w porównaniu z tradycyjnymi stałymi kodowaniami pozycyjnymi.
6. Uwaga: Grok-1 wykorzystuje 48 głów uwagi dla zapytań i 8 głów uwagi dla kluczy i wartości, z których każdy ma rozmiar 128.
7. Długość kontekstu: Model może przetwarzać ciągi o długości do 8 192 tokenów, wykorzystując precyzję bfloat16 do efektywnej obliczeń.

Wydajność i szczegóły wdrożeniowe

Grok-1 wykazał imponującą wydajność, przewyższając LLaMa 2 70B i Mixtral 8x7B z wynikiem MMLU 73%, prezentując swoją efektywność i dokładność w różnych testach.

Jednak należy zauważyć, że Grok-1 wymaga znacznych zasobów GPU ze względu na swój rozmiar. Aktualna implementacja w wydaniu open-source koncentruje się na walidacji poprawności modelu i wykorzystuje niewydajną implementację warstwy MoE, aby uniknąć potrzeby niestandardowych jąder.

Niemniej, model obsługuje sharding aktywacji i kwantyzację 8-bitową, co może zoptymalizować wydajność i zmniejszyć wymagania pamięciowe.

W godnym uwagi kroku xAI wydał Grok-1 na licencji Apache 2.0, udostępniając jego wagi i architekturę społeczności globalnej do użytku i wkładu.

Wydanie open-source obejmuje przykładowy kod repozytorium JAX, który pokazuje, jak załadować i uruchomić model Grok-1. Użytkownicy mogą pobrać wagi checkpoint za pomocą klienta torrent lub bezpośrednio za pośrednictwem HuggingFace Hub, ułatwiając dostęp do tego przełomowego modelu.

Przyszłość Mixture-of-Experts w modelach językowych

Ponieważ popyt na większe i bardziej zaawansowane modele językowe nadal rośnie, przyjęcie technik MoE prawdopodobnie zyska na impetzie. Trwające badania koncentrują się na rozwiązaniu pozostałych wyzwań, takich jak poprawa stabilności szkolenia, złagodzenie przeuczenia się podczas dostosowania i optymalizacja wymagań pamięciowych i komunikacyjnych.

Jednym z obiecujących kierunków jest eksploracja hierarchicznych architektur MoE, gdzie każdy ekspert składa się z wielu pod-ekspertów. Ten podejście może potencjalnie umożliwić jeszcze większą skalowalność i efektywność obliczeniową, przy zachowaniu wyrażalności dużych modeli.

Ponadto rozwój systemów sprzętowych i oprogramowania zoptymalizowanych dla modeli MoE jest aktywnym obszarem badań. Specjalistyczne przyspieszniki i ramy szkolenia rozproszonego zaprojektowane w celu efektywnej obsługi rzadkich i warunkowych wzorców obliczeniowych modeli MoE mogą dalej poprawić ich wydajność i skalowalność.

Dodatkowo, integracja technik MoE z innymi osiągnięciami w modelowaniu językowym, takimi jak mechanizmy uwagi rzadkiej, efektywne strategie tokenizacji i reprezentacje multimodalne, mogą prowadzić do jeszcze bardziej potężnych i wszechstronnych modeli językowych zdolnych do radzenia sobie z szerokim zakresem zadań.

Podsumowanie

Technika Mixture-of-Experts wyłoniła się jako potężne narzędzie w poszukiwaniu większych i bardziej zaawansowanych modeli językowych. Dzięki selektywnemu aktywowaniu ekspertów w oparciu o dane wejściowe, modele MoE oferują obiecujące rozwiązanie problemów obliczeniowych związanych ze skalowaniem gęstych modeli. Chociaż nadal istnieją wyzwania do pokonania, takie jak niestabilność szkolenia, przeuczenie się i wymagania pamięciowe, potencjalne korzyści modeli MoE w zakresie efektywności obliczeniowej, skalowalności i zrównoważoności środowiskowej sprawiają, że są one interesującym obszarem badań i rozwoju.

Ponieważ dziedzina przetwarzania języka naturalnego nadal posuwa się do przodu, przyjęcie technik MoE prawdopodobnie odegra kluczową rolę w umożliwieniu następnej generacji modeli językowych. Poprzez połączenie MoE z innymi osiągnięciami w architekturze modelu, technikach szkolenia i optymalizacji sprzętu, możemy spodziewać się jeszcze bardziej potężnych i wszechstronnych modeli językowych, które mogą prawdziwie zrozumieć i komunikować się z ludźmi w naturalny i bezproblemowy sposób.