Rozwój dużych modeli językowych (LLM) zbudowanych z modeli transformatorów przeznaczonych wyłącznie do dekodera odegrał kluczową rolę w transformacji domeny przetwarzania języka naturalnego (NLP), a także w rozwoju różnorodnych zastosowań głębokiego uczenia się, w tym uczenie się wzmacniania, analiza szeregów czasowych, przetwarzanie obrazu i wiele więcej. Jednak pomimo swojej skalowalności i dużej wydajności, LLM zbudowane z modeli transformatorów przeznaczonych wyłącznie do dekodera nadal borykają się z poważnymi niedociągnięciami. Chociaż ekspresyjny, mechanizm uwagi w LLM wywodzących się z transformatora wymaga dużych zasobów obliczeniowych zarówno podczas wnioskowania, jak i uczenia, co wymaga znacznej pamięci na długość sekwencji i kwadratowe FLOP. Te wysokie wymagania obliczeniowe ograniczają długość kontekstu modeli transformatorów, sprawiając, że zadania generowania autoregresyjnego są proporcjonalnie kosztowne w skali i utrudniają uczenie się na podstawie ciągłych strumieni danych oraz możliwość naprawdę nieograniczonego przetwarzania sekwencji.

W ostatnich czasach, Modele przestrzeni stanów (SSM) wykazały niezwykłe możliwości i wydajność, konkurując z modelami architektury transformatorów w testach porównawczych modelowania na dużą skalę, osiągając jednocześnie złożoność pamięci jako funkcję długości sekwencji i czasu liniowego. Co więcej, Mamba, niedawno wydany model przestrzeni stanów, wykazała się wyjątkową wydajnością w szeregu zadań związanych z modelowaniem języka i przetwarzaniem długich sekwencji. Jednocześnie modele Mixture of Expert (MoE) również wykazały imponującą wydajność, jednocześnie znacznie zmniejszając opóźnienia i koszty obliczeniowe wnioskowania, aczkolwiek kosztem większego zużycia pamięci. Opierając się na modelach Mamba i MoE, w tym artykule omówimy BlackMamba, nowatorską architekturę, która łączy model przestrzeni stanowej Mamba z modelami MoE, aby wykorzystać korzyści oferowane przez oba frameworki. Eksperymenty na platformie BlackMamba wykazały, że jest ona w stanie przewyższać istniejącą platformę Mamba i podstawy transformatorów zarówno pod względem uczenia FLOP, jak i wnioskowania. Wyjątkowa wydajność platformy BlackMamba pokazuje, że może ona skutecznie łączyć możliwości platform Mamba i MoE, oferując szybkie i ekonomiczne wnioskowanie z MoE z generowaniem złożoności liniowej z Mamby.

Celem tego artykułu jest szczegółowe omówienie frameworku BlackMamba. Badamy mechanizm, metodologię i architekturę frameworka, a także jego porównanie z najnowocześniejszymi frameworkami do generowania obrazów i wideo. Zacznijmy.

BlackMamba: wprowadzenie do MoE dla modeli przestrzeni stanów

Rozwój modeli wielkojęzycznych (LLM), szczególnie tych opartych na architekturach transformatorowych składających się wyłącznie z dekodera, w znaczący sposób wpłynął na Przetwarzanie języka naturalnego (NLP) dziedzinie i rozszerzony na różne zastosowania głębokiego uczenia się, w tym uczenie się przez wzmacnianie, analizę szeregów czasowych, przetwarzanie obrazu i nie tylko. Niemniej jednak, pomimo ich skalowalności i solidnej wydajności, te LLM oparte wyłącznie na dekoderze i transformatorach napotykają znaczące wyzwania. Mechanizm uwagi, kluczowa cecha oparta na transformatorze LLMs wymaga rozległych zasobów obliczeniowych zarówno do wnioskowania, jak i szkolenia. Wiąże się to z zapotrzebowaniem na pamięć, która rośnie wraz z długością sekwencji i operacjami obliczeniowymi (FLOP), które rosną kwadratowo. Tak intensywne potrzeby obliczeniowe ograniczają długość kontekstu modeli, podnoszą koszty zadań generowania autoregresji w miarę skalowania modelu i utrudniają zdolność modeli do efektywnego uczenia się na podstawie ciągłych strumieni danych lub sekwencji procesów o nieograniczonej długości. 

W ciągu ostatnich kilku lat poczyniono znaczne wysiłki, próbując przezwyciężyć te ograniczenia, i uwaga została przesunięta w stronę opracowania architektonicznych alternatyw dla kanonicznych modeli transformatorów o gęstej uwadze, przy czym najbardziej obiecującymi kandydatami są modele SSM i MoE. Kluczową korzyścią wynikającą z faworyzowania modeli przestrzeni stanów nad modelami architektury transformatorów jest liniowa złożoność obliczeniowa w odniesieniu do długości sekwencji wejściowej oferowana przez SSM w przeciwieństwie do złożoności kwadratowej oferowanej przez transformatory. Teoretycznie liniowa złożoność obliczeniowa w odniesieniu do długości sekwencji wejściowej umożliwia modelom przestrzeni stanów przetwarzanie większych sekwencji niż modele architektury transformatora dla danego budżetu FLOPS lub operacji zmiennoprzecinkowych na sekundę oraz renderowanie stałej generacji autoregresyjnej w obliczeniach bez pamięci podręcznej KV. Niedawno opracowane modele przestrzeni stanów, w tym Mamba, RetNet i kilka innych, wykazały skuteczne wnioskowanie i uczenie o długich sekwencjach, wraz z konkurencyjną wydajnością zadań modelowania językowego w przypadku transformatorów o podobnych właściwościach skalowania. Z drugiej strony, architektura oparta na mieszance modeli eksperckich zyskuje na popularności jako alternatywa dla gęstych transformatorów, ponieważ umożliwia znaczną redukcję wnioskowania i uczenia FLOP niezbędnych do osiągnięcia jakości porównywalnej z gęstym modelem. Modele MoE (Mixture of Experts) działają poprzez aktywację jedynie rzadkiego wyboru całkowitych parametrów podczas pojedynczego przejścia do przodu. Wykorzystują funkcję routingu, aby określić, którzy „eksperci” zostaną wezwani do działania w oparciu o dany kontekst. Podejście to oddziela koszt obliczeniowy wnioskowania od całkowitej liczby parametrów, umożliwiając lepszą wydajność w ramach ustalonego budżetu wnioskowania, aczkolwiek przy zwiększonej liczbie parametrów i większym zapotrzebowaniu na pamięć.

Ten postęp w architekturze oferuje znaczące korzyści w porównaniu z tradycyjnymi transformatorami i stanowi ekscytujący kierunek dalszego rozwoju. Zakładamy, że zintegrowanie tych ulepszeń w połączonym modelu Mamba-MoE mogłoby znacznie przyspieszyć możliwości modelowania języka i wydajność w porównaniu ze standardowymi modelami transformatorów. Przewidywane zalety architektury Mamba-MoE w porównaniu z tradycyjnym modelem gęstego transformatora obejmują:

Mamba: Osiąga liniową złożoność obliczeniową w stosunku do długości sekwencji wejściowej zarówno dla fazy uczenia, jak i wnioskowania. Umożliwia generowanie autoregresyjne w stałych ramach czasowych i przy stałym wykorzystaniu pamięci.

Ministerstwo Środowiska: Oferuje szybkość wnioskowania i wydajność obliczeniową szkolenia porównywalną z mniejszym, gęstym modelem bazowym, przy jednoczesnym zachowaniu poziomu jakości modelu, który może konkurować z modelem o równoważnej liczbie parametrów co wersja o większej gęstości.

Mając to na uwadze, należy stwierdzić, że modele architektury transformatorów są nadal najnowocześniejsze i wykazały stałą i niezwykle wysoką wydajność w zadaniach modelowania języka i zadań przetwarzania sekwencji. W swojej istocie architektura transformatora wykorzystuje samouważność, która wykonuje kwadratowe porównanie wszystkich do wszystkich podobieństw iloczynu skalarnego pomiędzy osadzeniem różnych tokenów w sekwencji i wykonuje liniowe odwzorowanie na wektor wyjściowy. Model transformatora składa się z bloków samouważności ułożonych pomiędzy blokami MLP lub wielowarstwowego perceptronu, które ponadto składają się z dwuwarstwowego MLP z daną funkcją aktywacji. 

BlackMamba: Architektura i metodologia

Modele przestrzeni stanów

Modele Przestrzeni Stanów należą do grupy modeli sekwencji o złożoności liniowej ze względu na długość ciągu wejściowego. Architektura modeli przestrzeni stanów jest bardziej zgodna z rekurencyjnymi sieciami neuronowymi i konwolucyjnymi sieciami neuronowymi niż architekturą opartą na uwadze i jest inspirowana ciągłym systemem dynamicznym, który odwzorowuje funkcję jednowymiarową poprzez ukrytą przestrzeń ukrytą. Liniowy układ dynamiczny sprawia, że ​​obliczenia równoległe są wydajne przy użyciu skanowania asocjacyjnego lub splotu. W praktycznych scenariuszach powtarzający się charakter modeli przestrzeni stanów był powodem, dla którego nadal nie można ich zastosować na wysoce równoległym sprzęcie AI, takim jak procesory graficzne. Jednak pojawienie się SSM, takich jak RWKV i Mamba wykorzystali jądra skanowania równoległego do wydajnego mapowania powtarzających się operacji na procesory graficzne, ułatwiając w ten sposób uczenie nowatorskich architektur z wydajnością porównywalną do tej osiąganej przez modele transformatorów. 

Nieodłączna złożoność kwadratowa w odniesieniu do długości sekwencji w transformatorach jest dobrze znanym ograniczeniem, które utrudnia rozumowanie i zrozumienie w bardzo długich kontekstach. Ostatnie innowacje wprowadziły pomysł wydłużania długości kontekstu, umożliwiając szkolenie transformatorów na wykonalną skalę przed zastosowaniem ich do znacznie dłuższych kontekstów podczas wnioskowania. Pomimo tych udoskonaleń proces wnioskowania w dalszym ciągu wymaga znacznej ilości zasobów obliczeniowych i pamięci, szczególnie do obsługi pamięci podręcznej typu klucz-wartość (KV), co czyni go przedsięwzięciem wymagającym dużych zasobów. Ostatnie wysiłki badawcze skupiły się na wzmocnieniu możliwości ekspresyjnych modeli przestrzeni stanów poprzez włączenie mechanizmów bramkowania zależnych od danych wejściowych, podobnych do macierzy Zapytanie, Klucz, Wartość (QKV) występujących w mechanizmach uwagi. 

Wysiłki te mają na celu zachowanie nieodłącznie liniowego postępu rekurencji w przestrzeni stanów, umożliwiając efektywne wykonanie poprzez splot lub proces selektywnego skanowania. Takie podejście znacznie zmniejsza różnice w wydajności transformatorów w zastosowaniach praktycznych. Wśród tych udoskonaleń Mamba wyróżnia się jako model przestrzeni stanów, który odzwierciedla cele wcześniejszych badań, wykazując imponujący poziom wydajności porównywalny z transformatorami w skalach do 2.8 miliarda parametrów. Osiąga się to poprzez zastosowanie bramkowania zależnego od danych wejściowych do wejść rekurencji modelu przestrzeni stanów (SSM), zapewniając jednocześnie wydajne obliczenia dzięki zastosowaniu specjalnie zaprojektowanych jąder selektywnego skanowania.

Mieszanka modeli eksperckich

Mieszanka modeli eksperckich (MoE) umożliwia rozdzielenie kosztu wnioskowania od całkowitej liczby parametrów poprzez selektywne aktywowanie parametrów podczas przebiegu w przód. Zamiast wykorzystywać wszystkie parametry, modele te kierują tokeny do konkretnych ekspertów Multilayer Perceptron (MLP). W idealnym przypadku każdy ekspert jest dostosowany do przetwarzania określonego rodzaju danych wejściowych, z mechanizmem routingu, zasadniczo zwartą siecią neuronową, określającym najbardziej odpowiedniego eksperta dla każdego tokena. Podejście to ma na celu zachowanie wszechstronnej mocy wyrazu modelu o równoważnej liczbie parametrów w gęstszej konfiguracji, ale przy znacznie zmniejszonych wymaganiach obliczeniowych. Zazwyczaj router mapuje warstwy liniowe od tokenów do indeksów eksperckich, przy czym każdy ekspert jest po prostu standardowym wielowarstwowym perceptronem transformatorowym. Jednak programiści nie odkryli jeszcze optymalnej metody szkolenia routera, ponieważ problemu przydziału ekspertów nie można różnicować, a modele mieszane często borykają się z problemem równoważenia obciążenia i stabilności treningu pomiędzy różnymi ekspertami w celu uzyskania wydajności sprzętu. 

Architektura

W swojej istocie BlackMamba wykorzystuje standardowy model transformatora składający się z przeplatanych bloków MLP i bloków uwagi dodawanych sekwencyjnie wzdłuż strumienia resztkowego. Obecnie większość modeli typu Mixture of Expert po prostu zastępuje wielowarstwowe bloki perceptronu trasowaną warstwą ekspercką. Z drugiej strony framework BlackMamba nie tylko zastępuje wielowarstwowy blok perceptronu w transformatorze trasowaną warstwą ekspercką, ale także zastępuje warstwę uwagi warstwą modelu przestrzeni stanu Mamba. Architekturę frameworku BlackMamba przedstawiono na poniższym rysunku. 

Szkolenie i zbiór danych

Model BlackMamba jest szkolony na ponad 300 miliardach tokenów w niestandardowym zestawie danych i wykorzystuje funkcję aktywacji SwiGLU dla doświadczonych perceptronów wielowarstwowych. Platforma szkoli się z udziałem 8 ekspertów, a liczba ta została uznana przez programistów za właściwą równowagę i kompromis pomiędzy zajmowaniem pamięci a kosztem wnioskowania modelu. Niestandardowy zestaw danych używany do uczenia platformy BlackMamba składa się z mieszanki już istniejących zestawów danych typu open source, w tym Starcoder, SlimPajama, Pile i innych. Poniższa tabela przedstawia wagi każdego zestawu danych użytego do szkolenia platformy BlackMamba. Ogółem w zbiorze danych znajduje się 1.8 biliona tokenów. 

BlackMamba: Wyniki

Aby zapewnić uczciwe porównanie Mamby i BlackMamby, programiści wytrenowali oba modele z tymi samymi parametrami treningowymi na tych samych danych treningowych. Framework BlackMamba jest w stanie przewyższyć zarówno modele Mamby, jak i modele transformatora, zapewniając identyczny rozmiar modelu przejścia w przód w czasie wnioskowania, a także szkoląc operacje zmiennoprzecinkowe na sekundę. Poniższy rysunek przedstawia czas potrzebny do wygenerowania sekwencji o danej długości w sposób autoregresyjny na podstawie początkowego jednoznakowego podpowiedzi w funkcji długości sekwencji. 

Co więcej, korzyści w zakresie opóźnień wynikające z połączenia modeli Expert i Mamba są połączone w frameworku BlackMamba, co skutkuje znacznie krótszymi czasami wnioskowania w porównaniu z modelami transformatorów, czystymi modelami Mamba i modelami MoE. Co więcej, przewaga wnioskowania platformy BlackMamba jest wprost proporcjonalna do długości sekwencji, dzięki czemu BlackMamba jest niezwykle skuteczna w generowaniu długich sekwencji. Idąc dalej, poniższy rysunek ilustruje liczbę tokenów przypisanych do modeli BlackMamba z odpowiednio 340 milionami i 640 milionami parametrów. Jak widać, większość warstw charakteryzuje się wysokim poziomem równowagi eksperckiej, co jest efektem ulepszonego algorytmu Sinkhorna zaimplementowanego w modelach BlackMamba. 

Poniższa tabela przedstawia wyniki oceny platformy BlackMamba w porównaniu z szeregiem wstępnie wytrenowanych modeli języków typu open source. Jak można zaobserwować, framework BlackMamba jest w stanie konkurować i osiągać lepsze wyniki niż większość frameworków we wszystkich wersjach bazowych. Co więcej, warto zauważyć, że modele, które przewyższają BlackMambę, mają znacznie większą liczbę parametrów, a różnica w wydajności jest minimalna, co wskazuje na możliwości frameworka BlackMamba przy mniejszych parametrach. 

Final Thoughts

W tym artykule omówiliśmy BlackMamba, nowatorską architekturę, która łączy model przestrzeni stanowej Mamba z mieszanką modeli eksperckich, aby czerpać korzyści oferowane przez oba te frameworki. Eksperymenty na platformie BlackMamba wykazały, że przewyższa ona istniejącą platformę Mamba i podstawy transformatorów zarówno pod względem uczenia FLOP, jak i wnioskowania. Wyjątkowa wydajność frameworka BlackMamba pokazuje, że jest on w stanie wyjątkowo dobrze dziedziczyć i łączyć możliwości frameworków Mamba i MoE, ponieważ łączy tanie i szybkie wnioskowanie z MoE z generowaniem o złożoności liniowej z Mamby. Rozmawialiśmy o tym, jak architektura frameworku BlackMamba jest w stanie przewyższyć dobrze wyszkolone modele dużego języka, istniejący framework Mamba i modele Mixture of Expert pod względem uczenia FLOP i kosztów wnioskowania. Co więcej, framework BlackMamba dziedziczy również generację FLOPów i ​​zredukowane szkolenie jednocześnie z obu modeli Mixture of Expert i frameworka Mamba.