Vzestup Mixture-of-Experts pro efektivní velké jazykové modely

Ve světě zpracování přirozeného jazyka (NLP) je snaha o stavbu větších a schopnějších jazykových modelů hnací silou mnoha nedávných pokroků. Nicméně, jak tyto modely rostou ve velikosti, výpočetní požadavky pro trénování a inferenci se stávají stále více náročnými, tlačící proti limitům dostupných hardwarových zdrojů.

Vstoupí Mixture-of-Experts (MoE), technika, která slibuje zmírnit tuto výpočetní zátěž a umožnit trénování větších a mocnějších jazykových modelů. Níže budeme diskutovat o MoE, prozkoumáme jeho původ, vnitřní fungování a jeho aplikace v transformátorových jazycových modelech.

Původ Mixture-of-Experts

Koncept Mixture-of-Experts (MoE) lze vysledovat zpět do počátku 90. let, kdy výzkumníci prozkoumali myšlenku podmíněného výpočtu, kde části neuronové sítě jsou selektivně aktivovány na základě vstupních dat. Jednou z průkopnických prací v tomto oboru byla práce “Adaptivní směs místních expertů” od Jacobs et al. v roce 1991, která navrhla supervidované učení pro soubor neuronových sítí, každá specializovaná na jinou oblast vstupního prostoru.

Jádro myšlenky za MoE spočívá v tom, že máme několik “expertních” sítí, každá zodpovědná za zpracování podmnožiny vstupních dat. Bránící mechanismus, obvykle neuronová síť sama, určuje, které experty by měly zpracovat daný vstup. Tento přístup umožňuje modelu alokovat své výpočetní zdroje efektivněji aktivací pouze relevantních expertů pro každý vstup, místo toho, aby využíval plnou kapacitu modelu pro každý vstup.

V průběhu let různí výzkumníci prozkoumali a rozvinuli myšlenku podmíněného výpočtu, vedoucí k vývoji technik, jako jsou hierarchické MoE, nízkorozměrové aproximace pro podmíněný výpočet a techniky pro odhadování gradientů prostřednictvím stochastických neuronů a hard-threshold aktivačních funkcí.

Mixture-of-Experts v Transformátorech

Mixture of Experts

Zatímco myšlenka MoE existuje již desetiletí, její aplikace na transformátorové jazykové modely je relativně nedávná. Transformátory, které se staly de facto standardem pro špičkové jazykové modely, se skládají z více vrstev, každá obsahující mechanismus sebe-pozornosti a hustou neuronovou síť (FFN).

Klíčová inovace spočívá v aplikaci MoE na transformátory, kdy se nahrazují husté FFN vrstvy řídce MoE vrstvami, každá se skládající z více expertních FFN a bránicího mechanismu. Bránicí mechanismus určuje, které experty by měly zpracovat každý vstupní token, umožňující modelu selektivně aktivovat pouze podmnožinu expertů pro daný vstupní řetězec.

Jedna z prvních prací, která demonstrovala potenciál MoE v transformátorech, byla práce “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer” od Shazeer et al. v roce 2017. Tato práce zavedla koncept řídce bránicí MoE vrstvy, která využívala bránicí mechanismus, který přidával řídkost a šum do procesu výběru expertů, zajišťující, že pouze podmnožina expertů byla aktivována pro každý vstup.

Od té doby několik dalších prací dále rozvinulo aplikaci MoE na transformátory, řešící výzvy, jako je nestabilita trénování, vyvážení zátěže a efektivní inferenci. Pozoruhodné příklady zahrnují Switch Transformer (Fedus et al., 2021), ST-MoE (Zoph et al., 2022) a GLaM (Du et al., 2022).

Výhody Mixture-of-Experts pro Jazykové Modely

Hlavní výhodou využití MoE v jazycových modelech je schopnost škálovat velikost modelu a současně udržet relativně konstantní výpočetní náklady během inferenci. Selektivně aktivací pouze podmnožiny expertů pro každý vstupní token mohou modely MoE dosáhnout expresivní síly mnohem větších hustých modelů, zatímco vyžadují podstatně méně výpočtů.

Například uvažujme jazykový model s hustou FFN vrstvou o 7 miliardách parametrů. Pokud nahradíme tuto vrstvu MoE vrstvou se osmi experty, každým s 7 miliardami parametrů, celkový počet parametrů se zvýší na 56 miliard. Nicméně během inferenci, pokud aktivujeme pouze dva experty na token, výpočetní náklady jsou ekvivalentní 14miliardovému hustému modelu, protože se provádějí dvě 7miliardové matice násobení.

Tato výpočetní efektivita během inferenci je zvláště cenná v nasazeních, kde jsou zdroje omezené, jako jsou mobilní zařízení nebo edge computery. Kromě toho snížené výpočetní požadavky během trénování mohou vést k podstatným úsporám energie a nižší uhlíkové stopě, což je v souladu s rostoucím důrazem na udržitelnou AI.

Výzvy a Úvahy

Zatímco modely MoE nabízejí přesvědčivé výhody, jejich přijetí a nasazení také přinášejí několik výzev a úvah:

1. Nestabilita Trénování: Modely MoE jsou známé tím, že jsou více náchylné k nestabilitě trénování ve srovnání s jejich hustými protějšky. Tento problém vzniká z řídké a podmíněné povahy aktivace expertů, což může vést k problémům s propagací gradientů a konvergencí. Techniky, jako je router z-loss (Zoph et al., 2022), byly navrženy k zmírnění těchto nestabilit, ale další výzkum je stále potřebný.
2. Finetuning a Přeučení: Modely MoE tendují k přeučení se během finetuningu, zvláště když má downstream úkol relativně malou datovou sadu. Toto chování je připsáno zvýšené kapacitě a řídkosti modelů MoE, které mohou vést k přehnané specializaci na trénovací data. Péčlivé regularizace a strategie finetuningu jsou vyžadovány k zmírnění tohoto problému.
3. Požadavky na Paměť: Zatímco modely MoE mohou snižovat výpočetní náklady během inferenci, často mají vyšší požadavky na paměť ve srovnání s hustými modely podobné velikosti. To je způsobeno tím, že všechny váhy expertů musí být načteny do paměti, i když je pouze podmnožina aktivována pro každý vstup. Omezení paměti mohou omezit škálovatelnost modelů MoE na zařízení s omezenými zdroji.
4. Vyvážení Zátěže: K dosažení optimální výpočetní efektivnosti je důležité vyvážit zátěž napříč experty, zajišťující, že žádný jediný expert není přetížen, zatímco ostatní zůstávají nevyužití. Toto vyvážení zátěže je obvykle dosaženo pomocnými ztrátami během trénování a pečlivým naladěním kapacitního faktoru, který určuje maximální počet tokenů, které lze přiřadit každému expertovi.
5. Komunikační Přepětí: V distribuovaných scénářích trénování a inferenci mohou modely MoE zavést další komunikační přepětí kvůli potřebě výměny aktivace a gradientní informace napříč experty umístěnými na různých zařízeních nebo akcelerátorech. Efektivní komunikační strategie a hardwarově vědomý design modelu jsou nezbytné k zmírnění tohoto přepětí.

Přes tyto výzvy potenciální výhody modelů MoE v umožnění větších a schopnějších jazycových modelů vedly k významným výzkumným úsilím na řešení a zmírnění těchto problémů.

Příklad: Mixtral 8x7B a GLaM

K ilustraci praktické aplikace MoE v jazycových modelech uvažujme dva pozoruhodné příklady: Mixtral 8x7B a GLaM.

Mixtral 8x7B je variantou MoE modelu Mistral, vyvinutého společností Anthropic. Skládá se z osmi expertů, každého s 7 miliardami parametrů, což vede k celkovému počtu 56 miliard parametrů. Nicméně během inferenci jsou aktivovány pouze dva experty na token, efektivní výpočetní náklady jsou ekvivalentní 14miliardovému hustému modelu.

Mixtral 8x7B prokázal působivý výkon, překonávající 70miliardový model Llama, zatímco nabízí mnohem rychlejší časy inferenci. Verze instruktivně upravená Mixtral 8x7B, nazvaná Mixtral-8x7B-Instruct-v0.1, byla také vydána, dále zvyšující jeho schopnosti při sledování přirozených jazykových instrukcí.

Dalším pozoruhodným příkladem je GLaM (Google (GOOGL ) Language Model), velký MoE model vyvinutý společností Google. GLaM využívá dekodér-only transformátorovou architekturu a byl trénován na masivní 1,6bilionové tokenové datové sadě. Model dosahuje působivého výkonu v few-shot a one-shot evaluacích, odpovídající kvalitě GPT-3, zatímco využívá pouze jednu třetinu energie potřebné pro trénování GPT-3.

Úspěch GLaM lze připsat jeho efektivní MoE architektuře, která umožnila trénování modelu s obrovským množstvím parametrů, zatímco udržovala rozumné výpočetní požadavky. Model také prokázal potenciál modelů MoE být více energeticky efektivními a ekologicky udržitelnými ve srovnání s jejich hustými protějšky.

Architektura Grok-1

GROK MIXTURE OF EXPERT

Grok-1 je transformátorový MoE model s unikátní architekturou navrženou pro maximalizaci efektivity a výkonu. Podívejme se na klíčové specifikace:

1. Parametry: S ohromujícím 314 miliardami parametrů je Grok-1 největším otevřeným LLM do dneška. Díky architektuře MoE jsou však pouze 25% váh (přibližně 86 miliard parametrů) aktivní v jeden okamžik, zvyšující tak zpracovatelské schopnosti.
2. Architektura: Grok-1 využívá architekturu Mixture-of-8-Experts, kde je každý token zpracován dvěma experty během inferenci.
3. Vrstvy: Model se skládá z 64 transformátorových vrstev, každá obsahující mechanismus sebe-pozornosti a hustou blokovou neuronovou síť.
4. Tokenizace: Grok-1 využívá SentencePiece tokenizátor s velikostí slovníku 131 072 tokenů.
5. Vložené a Pozicionální Kódování: Model má 6 144-rozměrné vložené a využívá rotační pozicionální kódování, umožňující dynamičtější interpretaci dat ve srovnání s tradičním pevným pozicionálním kódováním.
6. Pozornost: Grok-1 využívá 48 hlav pozornosti pro dotazy a 8 hlav pozornosti pro klíče a hodnoty, každá o velikosti 128.
7. Délka Kontextu: Model může zpracovat sekvence až 8 192 tokenů, využívající bfloat16 přesnost pro efektivní výpočty.

Výkon a Implementační Detaily

Grok-1 prokázal působivý výkon, překonávající LLaMa 2 70B a Mixtral 8x7B s MMLU skóre 73%, demonstrující jeho efektivitu a přesnost napříč různými testy.

Je však důležité poznamenat, že Grok-1 vyžaduje značné GPU zdroje kvůli své velikosti. Současná implementace v otevřeném vydání se zaměřuje na ověření správnosti modelu a využívá neefektivní implementaci MoE vrstvy, aby se zabránilo potřebě vlastních jader.

Nicméně model podporuje aktivaci shardingu a 8bitovou kvantizaci, které mohou optimalizovat výkon a snížit požadavky na paměť.

V pozoruhodném kroku xAI vydal Grok-1 pod licencí Apache 2.0, zpřístupňující jeho váhy a architekturu globální komunitě pro použití a příspěvky.

Otevřené vydání zahrnuje JAX ukázkový repozitář kódu, který demonstruje, jak načíst a spustit model Grok-1. Uživatelé mohou stáhnout checkpoint váhy pomocí klienta torrentu nebo přímo prostřednictvím HuggingFace Hub, usnadňující přístup k tomuto průlomovému modelu.

Budoucnost Mixture-of-Experts v Jazycových Modelech

Jak poptávka po větších a schopnějších jazycových modelech pokračuje, je pravděpodobné, že přijetí technik MoE získá další impuls. Pokračující výzkumná úsilí se zaměřují na řešení zbývajících výzev, jako je zlepšení stability trénování, zmírnění přeučení během finetuningu a optimalizace požadavků na paměť a komunikaci.

Jedním slibným směrem je prozkoumání hierarchických MoE architektur, kde je každý expert sám složen z více sub-expertů. Tento přístup by mohl potenciálně umožnit ještě větší škálovatelnost a výpočetní efektivitu, zatímco udržuje expresivní sílu velkých modelů.

Kromě toho vývoj hardwaru a softwaru optimalizovaného pro modely MoE je aktivní oblastí výzkumu. Specializované akcelerátory a distribuované trénovací rámce navržené pro efektivní zpracování řídkých a podmíněných výpočtů modelů MoE by mohly dále zvýšit jejich výkon a škálovatelnost.

Dále integrace technik MoE s jinými pokroky v jazycovém modelování, jako jsou sparse pozornostní mechanismy, efektivní tokenizace a multimodální reprezentace, by mohla vést k ještě mocnějším a všestrannějším jazycovým modelům schopným řešit širokou škálu úkolů.

Závěr

Technika Mixture-of-Experts se stala mocným nástrojem v hledání větších a schopnějších jazycových modelů. Selektivně aktivací expertů na základě vstupních dat nabízejí modely MoE slibné řešení výpočetních výzev spojených se škálováním hustých modelů. Zatímco existují ještě výzvy, které je třeba překonat, jako je nestabilita trénování, přeučení a požadavky na paměť, potenciální výhody modelů MoE v terms of výpočetní efektivitě, škálovatelnosti a environmentální udržitelnosti je činí zajímavou oblastí výzkumu a vývoje.

Jak oblast zpracování přirozeného jazyka pokračuje v tlačení hranic toho, co je možné, je pravděpodobné, že přijetí technik MoE bude hrát klíčovou roli v umožnění další generace jazycových modelů. Kombinací MoE s jinými pokroky v modelové architektuře, trénovacích technikách a hardwarové optimalizaci můžeme očekávat ještě mocnější a všestrannější jazycové modely, které mohou skutečně rozumět a komunikovat s lidmi přirozeným a bezproblémovým způsobem.