Vývoj velkých jazykových modelů (LLM) vybudovaných z modelů transformátorů pouze s dekodérem sehrál klíčovou roli při transformaci domény zpracování přirozeného jazyka (NLP) a také v rozvoji různých aplikací hlubokého učení, včetně posilování učení, analýza časových řad, zpracování obrazu a mnoho dalšího. Navzdory své škálovatelnosti a vysokému výkonu však LLM postavené z modelů transformátorů pouze s dekodérem stále čelí významným nedostatkům. Ačkoli je mechanismus pozornosti v LLM odvozených z transformátoru expresivní, vyžaduje vysoké výpočetní zdroje během inference i tréninku, což vyžaduje značnou paměť pro délku sekvence a kvadratické FLOP. Tento vysoký výpočetní požadavek omezuje kontextovou délku modelů transformátorů, čímž se úlohy autoregresního generování úměrně prodražují s měřítkem a brání učení z nepřetržitých datových toků a schopnosti skutečně neomezeného zpracování sekvencí.

V nedávné době, Modely státního prostoru (SSM) prokázaly pozoruhodné schopnosti a výkon a konkurovaly modelům architektury transformátoru ve velkých modelovacích benchmarcích, přičemž dosáhly složitosti paměti jako funkce délky sekvence a lineárního času. Kromě toho Mamba, nedávno vydaný State Space Model, prokázal vynikající výkon v řadě úloh jazykového modelování a zpracování dlouhých sekvencí. Současně modely Mixture of Expert (MoE) také ukázaly působivý výkon a zároveň výrazně snížily latenci a výpočetní náklady odvozené, i když na úkor větší paměti. V návaznosti na modely Mamba a MoE bude tento článek diskutovat o BlackMambě, nové architektuře, která kombinuje model státního prostoru Mamba s modely MoE, aby se využily výhody nabízené oběma frameworky. Experimenty na BlackMamba prokázaly její schopnost překonat stávající rámec Mamba a základní linie transformátoru jak při trénování FLOP, tak při vyvozování. Výjimečný výkon rámce BlackMamba ukazuje, že dokáže efektivně kombinovat schopnosti rámců Mamba a MoE a nabízí rychlé a cenově výhodné odvození z MoE s generováním lineární složitosti z Mamba.

Tento článek si klade za cíl pokrýt do hloubky rámec BlackMamba. Prozkoumáme mechanismus, metodologii a architekturu rámce spolu s jeho srovnáním s nejmodernějšími rámci pro generování obrázků a videa. Začněme.

BlackMamba: Úvod do MŽP pro modely státního prostoru

Vývoj velkých jazykových modelů (LLM), zejména těch, které jsou založeny na transformátorových architekturách pouze pro dekodéry, výrazně ovlivnil Zpracování přirozeného jazyka (NLP) a rozšířil se do různých aplikací hlubokého učení, včetně posilovacího učení, analýzy časových řad, zpracování obrazu a dalších. Navzdory své škálovatelnosti a robustnímu výkonu se však tyto LLM založené pouze na dekodérech potýkají se značnými problémy. Mechanismus pozornosti, klíčová vlastnost transformátoru LLMs, vyžaduje rozsáhlé výpočetní zdroje pro vyvozování i školení. To zahrnuje potřebu paměti, která roste s délkou sekvence a výpočetními operacemi (FLOP), které rostou kvadraticky. Takové intenzivní výpočetní potřeby omezují délku kontextu modelů, zvyšují náklady na úlohy autoregresního generování, jak se model mění, a brání schopnosti modelů učit se z kontinuálních datových toků nebo efektivně zpracovávat sekvence neomezené délky. 

V posledních několika letech bylo vynaloženo značné úsilí ve snaze překonat tato omezení a pozornost byla přesunuta k navrhování architektonických alternativ ke kanonickým modelům transformátorů s hustou pozorností, přičemž modely SSM a MoE jsou nejslibnějšími kandidátskými architekturami. Klíčovou výhodou upřednostnění modelů stavového prostoru před modely architektury transformátorů je lineární výpočetní složitost s ohledem na délku vstupní sekvence, kterou nabízejí SSM, na rozdíl od kvadratické složitosti nabízené transformátory. Teoreticky lineární výpočetní složitost s ohledem na délku vstupní sekvence umožňuje modelům stavového prostoru zpracovávat větší sekvence než modely architektury transformátoru pro daný rozpočet FLOPS nebo operace s plovoucí desetinnou čárkou a vykreslovat konstantu autoregresního generování při výpočtu bez KV cache. Nedávno vyvinuté modely státního prostoru, včetně Mamba, RetNet a několika dalších, prokázaly efektivní dlouhodobou inferenci a školení spolu s konkurenčním výkonem úlohy modelování jazyka pro transformátory s podobnými vlastnostmi škálování. Na druhé straně si architektura modelů Mixture of Expert získává na popularitě jako alternativa k hustým transformátorům, protože umožňuje výrazné snížení inference a trénování FLOP nezbytných pro dosažení srovnatelné kvality s hustým modelem. Modely MoE (Mixture of Experts) fungují tak, že během jednoho průchodu vpřed aktivují pouze řídký výběr celkových parametrů. Využívají směrovací funkci k určení, kteří „experti“ jsou na základě daného kontextu povoláni do akce. Tento přístup vytváří oddělení mezi výpočetními náklady na odvození a celkovým počtem parametrů, což umožňuje vyšší výkon v rámci fixního rozpočtu na odvození, i když se zvýšeným počtem parametrů a většími požadavky na paměť.

Tento pokrok v architektuře nabízí oproti tradičním transformátorům značné výhody a představuje vzrušující směr pro další vývoj. Předpokládáme, že integrace těchto vylepšení do kombinovaného modelu Mamba-MoE by mohla výrazně urychlit možnosti jazykového modelování a efektivitu nad rámec standardních modelů transformátorů. Mezi očekávané výhody architektury Mamba-MoE ve srovnání s tradičním modelem hustého transformátoru patří:

mamba: Dosahuje lineární výpočetní složitosti vzhledem k délce vstupní sekvence pro trénovací i inferenční fázi. Umožňuje, aby k autoregresivnímu generování docházelo v konstantním časovém rámci as konstantním využitím paměti.

Vočko: Nabízí rychlost odvození a trénovací výpočetní efektivitu srovnatelnou s menším, hustým základním modelem při zachování úrovně kvality modelu, která konkuruje modelu s ekvivalentním počtem parametrů jako hustší verze.

S tím, co bylo řečeno, je nezbytné uvést, že modely architektury transformátorů jsou stále nejmodernější a prokázaly konzistentní a pozoruhodně silný výkon v úlohách jazykového modelování a úlohách zpracování sekvencí. Ve svém jádru využívá architektura transformátoru sebepozornost, která provádí kvadratické srovnání podobností bodového součinu mezi vloženími různých tokenů v sekvenci a provádí lineární mapu k výstupnímu vektoru. Model transformátoru se skládá z bloků samopozornosti naskládaných mezi bloky MLP nebo Multi-Layer Perceptron, které se dále skládají z dvouvrstvého MLP s danou aktivační funkcí. 

BlackMamba: Architektura a metodologie

Modely státního prostoru

Stavové modely patří do skupiny sekvenčních modelů s lineární složitostí vzhledem k délce vstupní sekvence. Architektura State Space Models je spíše v souladu s rekurentními neuronovými sítěmi a konvolučními neuronovými sítěmi než s architekturou založenou na pozornosti a je inspirována kontinuálním dynamickým systémem, který mapuje 1-rozměrnou funkci prostřednictvím implicitního latentního prostoru. Lineární dynamický systém zefektivňuje paralelní výpočty pomocí asociativního nebo konvolučního skenování. V praktických scénářích byla opakující se povaha State Space Models důvodem, proč je stále třeba přijmout na vysoce paralelní hardware AI, jako jsou GPU. Nicméně, vznik SSM jako RWKV a Mamba použili jádra paralelního skenování k efektivnímu mapování opakujících se operací na GPU, čímž usnadnili trénování nových architektur s účinností srovnatelnou s těmi, které dosahují modely transformátorů. 

Inherentní kvadratická složitost ve vztahu k délce sekvence v transformátorech je dobře známým omezením, které brání uvažování a porozumění ve velmi dlouhých kontextech. Nedávné inovace zavedly myšlenku prodloužení délky kontextu, což umožňuje, aby byly transformátory trénovány v proveditelném měřítku před tím, než budou aplikovány na mnohem delší kontexty během inference. Navzdory těmto pokrokům proces vyvozování stále vyžaduje značné množství výpočetních zdrojů a paměti, zejména pro údržbu mezipaměti klíče a hodnoty (KV), což z něj činí úsilí náročné na zdroje. Nedávné výzkumné úsilí se zaměřilo na zlepšení vyjadřovacích schopností modelů stavového prostoru začleněním vstupně závislých hradlových mechanismů, podobných maticím dotazů, klíčů, hodnot (QKV), které se nacházejí v mechanismech pozornosti. 

Cílem těchto snah je zachovat přirozeně lineární průběh rekurze stavového prostoru, což umožňuje efektivní provádění buď pomocí konvoluce nebo procesu selektivního skenování. Tento přístup výrazně zužuje disparitu výkonu s transformátory v praktických aplikacích. Mezi těmito vylepšeními vyniká Mamba jako model stavového prostoru, který odráží cíle předchozího výzkumu a vykazuje působivé úrovně výkonu srovnatelné s transformátory v měřítku až 2.8 miliardy parametrů. Dosahuje toho použitím vstupně závislého hradlování na vstupy rekurze stavového modelu (SSM), a to při současném zajištění efektivních výpočtů díky použití zakázkových selektivních skenovacích jader.

Směs expertních modelů

Mixture of Expert (MoE) modely dosahují oddělení mezi inferenčními náklady a celkovým počtem parametrů selektivní aktivací parametrů během dopředného průchodu. Namísto použití všech parametrů tyto modely směrují tokeny konkrétním odborníkům na vícevrstvý perceptron (MLP). V ideálním případě je každý expert uzpůsoben tak, aby zpracovával určitý typ vstupu, se směrovacím mechanismem, v podstatě kompaktní neuronovou sítí, který určuje nejvhodnějšího experta pro každý token. Tento přístup si klade za cíl zachovat komplexní vypovídací schopnost modelu s ekvivalentním počtem parametrů v hustší konfiguraci, ale se značně sníženými výpočetními nároky. Směrovač je obvykle mapováním lineárních vrstev od tokenů po expertní indexy, přičemž každý expert je jednoduše standardní transformátor Multilayer Perceptron. Vývojáři však ještě musí přijít na optimální metodu školení pro směrovač, protože problém s přiřazením expertů je nerozlišitelný a modely Mixture of Expert se často potýkají s vyrovnáváním zátěže a stabilitou školení mezi různými odborníky pro efektivitu hardwaru. 

Architektura

BlackMamba ve svém jádru využívá standardní model transformátoru sestávající z prokládaných bloků MLP a bloků pozornosti přidávaných postupně podél zbytkového proudu. Nyní většina modelů Mixture of Expert jednoduše nahrazuje vícevrstvé perceptronové bloky směrovanou expertní vrstvou. Na druhou stranu framework BlackMamba nejen nahrazuje vícevrstvý perceptronový blok v transformátoru směrovanou expertní vrstvou, ale také nahrazuje vrstvu pozornosti vrstvou Mamba State Space Model. Architektura frameworku BlackMamba je znázorněna na následujícím obrázku. 

Školení a datová sada

Model BlackMamba je trénován na více než 300 miliardách tokenů na vlastní datové sadě a používá aktivační funkci SwiGLU pro expertní vícevrstvé perceptrony. Rámec trénuje s 8 odborníky, což je počet, který vývojáři shledali jako správný poměr a kompromis mezi paměťovou stopou a odvozenými náklady modelu. Vlastní datová sada používaná k trénování rámce BlackMamba se skládá ze směsi již existujících datových sad s otevřeným zdrojovým kódem včetně Starcoder, SlimPajama, Pile a dalších. Následující tabulka ukazuje váhy každého datového souboru použitého pro trénování rámce BlackMamba. Celkově je v datové sadě 1.8 bilionu tokenů. 

BlackMamba : Výsledky

Aby bylo zajištěno spravedlivé srovnání mezi Mamba a BlackMamba, vývojáři trénovali oba modely se stejnými tréninkovými parametry na stejných tréninkových datech. Rámec BlackMamba je schopen překonat modely Mamba i transformátory pro identickou velikost modelu dopředného průchodu v inferenčním čase, stejně jako trénování operací s pohyblivou řádovou čárkou za sekundu. Následující obrázek ukazuje čas potřebný k vygenerování sekvence dané délky autoregresivně z počáteční výzvy s jedním tokenem jako funkci délky sekvence. 

Kromě toho jsou výhody latence obou modelů Mixture of Expert a Mamba kombinovány v rámci BlackMamba, což vede k výrazně rychlejším inferencím ve srovnání s modely transformátorů, čistými modely Mamba a modely MoE. Kromě toho je inferenční výhoda rámce BlackMamba přímo úměrná délkám sekvencí, díky čemuž je BlackMamba extrémně efektivní při generování dlouhých sekvencí. Následující obrázek ukazuje počet tokenů přiřazených k modelům BlackMamba s 340 miliony a 640 miliony parametrů. Jak je vidět, většina vrstev vykazuje vysokou úroveň expertní rovnováhy v důsledku vylepšeného algoritmu Sinkhorn implementovaného modely BlackMamba. 

Následující tabulka obsahuje skóre hodnocení rámce BlackMamba ve srovnání s řadou předtrénovaných jazykových modelů s otevřeným zdrojovým kódem. Jak lze pozorovat, framework BlackMamba je schopen konkurovat a překonat většinu rámců napříč všemi základními liniemi. Dále stojí za zmínku, že modely, které překonávají BlackMambu, mají podstatně vyšší počet parametrů a mezera ve výkonu je minimální, což ukazuje na schopnost rámce BlackMamba s menšími parametry. 

Závěrečné myšlenky

V tomto článku jsme hovořili o BlackMamba, nové architektuře, která kombinuje model státního prostoru Mamba s modely Mixture of Expert, aby těžila z výhod nabízených oběma těmito frameworky. Experimenty na BlackMamba prokázaly, že překonává stávající rámec Mamba a základní linie transformátoru jak v tréninkových FLOPech, tak v inferencích. Výjimečný výkon rámce BlackMamba demonstruje, že je schopen výjimečně zdědit a kombinovat schopnosti rámců Mamba a MoE, protože kombinuje levné a rychlé odvození z MoE s generováním lineární složitosti z Mamba. Mluvili jsme o tom, jak je architektura rámce BlackMamba schopna překonat silné trénované velké jazykové modely, existující rámec Mamba a modely Mixture of Expert, pokud jde o trénink FLOP a náklady na odvození. Kromě toho framework BlackMamba také zdědí generační FLOP a redukovaný trénink z obou modelů Mixture of Expert a frameworku Mamba současně.