Rozwój dużych modeli językowych (LLM) opartych na modelach transformatorowych tylko z dekodera odegrał kluczową rolę w przekształceniu dziedziny przetwarzania języka naturalnego (NLP), a także w rozwoju różnych aplikacji głębokiego uczenia, w tym uczenia wzmocnionego, analizy szeregów czasowych, przetwarzania obrazów i wielu innych. Niemniej jednak, pomimo ich skalowalności i silnej wydajności, LLM oparte na modelach transformatorowych tylko z dekodera nadal napotykają znaczące ograniczenia. Chociaż są one wyraziste, mechanizm uwagi w modelach transformatorowych wymaga dużych zasobów obliczeniowych podczas inferencji i treningu, wymagając znacznej ilości pamięci dla długości sekwencji i operacji (FLOPs) zwiększających się kwadratowo. To duże zapotrzebowanie na obliczenia ogranicza długość kontekstu modeli, zwiększa koszty zadań generowania autoregresyjnego wraz ze skalą modelu i utrudnia uczenie się z ciągłych strumieni danych i przetwarzanie sekwencji o nieograniczonej długości.

W ostatnich czasach Modele Przestrzeni Stanu (SSM) wykazały się znaczącymi możliwościami i osiągnięciami, konkurując z modelami architektury transformatorowej w dużych benchmarkach modelowania, osiągając złożoność pamięciową jako funkcję długości sekwencji i liniowy czas. Co więcej, Mamba, niedawno wydany Model Przestrzeni Stanu, wykazał się znakomitą wydajnością w zakresie modelowania języka i przetwarzania długich sekwencji. Równocześnie, Modele Mieszanki Ekspertów (MoE) również wykazały się imponującą wydajnością, znacznie redukując opóźnienia i koszty obliczeniowe inferencji, chociaż za cenę większego zużycia pamięci. Budując na modelach Mamba i MoE, ten artykuł omawia BlackMamba, nową architekturę, która łączy Model Przestrzeni Stanu Mamba z modelami MoE, aby wykorzystać korzyści oferowane przez obie ramy. Eksperymenty na BlackMamba wykazały jego zdolność do przewyższenia istniejącej ramy Mamba i modeli transformatorowych w treningu FLOPs i inferencji. Wyjątkowa wydajność ramy BlackMamba pokazuje, że może ona skutecznie łączyć zdolności ram Mamba i MoE, oferując szybką i efektywną inferencję z MoE oraz generację o liniowej złożoności z Mamba.

Ten artykuł ma na celu omówienie ramy BlackMamba w szczegółach. Eksplorujemy mechanizm, metodologię i architekturę ramy, wraz z porównaniem do ram generacji obrazów i wideo. Zaczynajmy.

BlackMamba: Wprowadzenie do MoE dla Modeli Przestrzeni Stanu

Postęp dużych modeli językowych (LLM), szczególnie tych opartych na architekturze transformatorowej tylko z dekodera, znacząco wpłynął na dziedzinę przetwarzania języka naturalnego (NLP) i rozwinął się w różnych aplikacjach głębokiego uczenia, w tym uczeniu wzmocnionym, analizie szeregów czasowych, przetwarzaniu obrazów i poza. Niemniej jednak, pomimo ich skalowalności i silnej wydajności, te modele LLM oparte na architekturze transformatorowej tylko z dekodera napotykają znaczące wyzwania. Mechanizm uwagi, kluczowa cecha modeli transformatorowych, wymaga dużych zasobów obliczeniowych podczas inferencji i treningu. To duże zapotrzebowanie na obliczenia ogranicza długość kontekstu modeli, zwiększa koszty zadań generowania autoregresyjnego wraz ze skalą modelu i utrudnia uczenie się z ciągłych strumieni danych lub przetwarzanie sekwencji o nieograniczonej długości.

W ostatnich latach dokonano znaczących wysiłków, aby przezwyciężyć te ograniczenia, a uwaga została skierowana na opracowanie alternatywnych architektur dla gęstych modeli transformatorowych, z Modelem Przestrzeni Stanu i Modelem Mieszanki Ekspertów jako najbardziej obiecującymi kandydatami. Kluczową korzyścią płynącą z faworyzowania Modeli Przestrzeni Stanu nad modelami architektury transformatorowej jest liniowa złożoność obliczeniowa w odniesieniu do długości sekwencji wejściowej oferowana przez Modele Przestrzeni Stanu, 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 Stanu przetwarzanie dłuższych sekwencji niż modele architektury transformatorowej dla danego budżetu FLOPS lub operacji zmiennoprzecinkowych na sekundę, oraz renderowanie zadań generowania autoregresyjnego stałych w czasie bez pamięci podręcznej KV. Ostatnio opracowane Modele Przestrzeni Stanu, w tym RWKV i Mamba, wykazały się wydajnymi długimi sekwencjami inferencji i treningu, wraz z konkurencyjną wydajnością zadań modelowania języka w porównaniu z transformatorami o podobnych właściwościach skalowania. Z drugiej strony, Modele Mieszanki Ekspertów zyskują na popularności jako alternatywa dla gęstych transformatorów, ponieważ umożliwiają znaczącą redukcję opóźnień i kosztów obliczeniowych inferencji, chociaż za cenę większego zużycia pamięci.

Ten postęp w architekturze oferuje znaczące korzyści w porównaniu z tradycyjnymi transformatorami i reprezentuje ekscytujący kierunek dalszego rozwoju. Uważamy, że integrowanie tych ulepszeń w połączonym modelu Mamba-MoE może znacząco przyspieszyć możliwości modelowania języka i efektywność ponad te oferowane przez standardowe modele transformatorowe. Oczekiwane korzyści architektury Mamba-MoE w porównaniu z tradycyjnym gęstym modelem transformatorowym obejmują:

Mamba: Osiąga liniową złożoność obliczeniową w odniesieniu do długości sekwencji wejściowej dla obu faz treningu i inferencji. Umożliwia generowanie autoregresyjne w stałym czasie i z użyciem stałej ilości pamięci.

MoE: Oferuje prędkość inferencji i efektywność obliczeniową porównywalną do mniejszego, gęstszego modelu bazowego, utrzymując poziom jakości modelu, który rywalizuje z modelem o równoważnej liczbie parametrów w wersji gęstszej.

Ponadto, pomimo że modele architektury transformatorowej nadal są najlepsze i wykazały się silną, konsekwentną wydajnością w zadaniach modelowania języka i przetwarzania sekwencji, ich rdzeń opiera się na uwadze, która wykonuje kwadratową porównanie wszystkich tokenów w sekwencji i wykonuje liniową mapę na wektor wyjściowy. Model transformatorowy składa się z bloków uwagi ułożonych między blokami wielowarstwowego perceptronu, które składają się z dwuwarstwowego perceptronu z daną funkcją aktywacji.

BlackMamba: Architektura i Metodologia

Modele Przestrzeni Stanu

Modele Przestrzeni Stanu należą do grupy modeli sekwencji o liniowej złożoności w odniesieniu do długości sekwencji wejściowej. Architektura Modeli Przestrzeni Stanu bardziej przypomina sieci neuronowe rekurencyjne i konwolucyjne niż architekturę opartą na uwadze i jest inspirowana przez ciągły system dynamiczny, który mapuje funkcję jednowymiarową przez niejawne przestrzeń latentną. System dynamiczny liniowy umożliwia wydajne obliczenia równoległe za pomocą skanu stowarzyszeniowego lub konwolucyjnego. W praktyce, natura rekurencyjna Modeli Przestrzeni Stanu była powodem, dla którego nie zostały one jeszcze w pełni zaadoptowane na wysoko-paralelnym sprzęcie AI, takim jak GPU. Niemniej jednak, pojawienie się Modeli Przestrzeni Stanu, takich jak RWKV i Mamba, wykorzystało jądra skanowania równoległego, aby mapować operacje rekurencyjne wydajnie na GPU, umożliwiając trening nowych architektur z wydajnością porównywalną do tej osiąganej przez modele transformatorowe.

Wewnętrzna złożoność kwadratowa w odniesieniu do długości sekwencji w transformatorach jest dobrze znanym ograniczeniem, które utrudnia rozumienie i wnioskowanie na bardzo długich kontekstach. Ostatnie innowacje wprowadziły ideę rozszerzania długości kontekstu, umożliwiając transformatorom trening na przystępnej skali przed zastosowaniem ich do znacznie dłuższych kontekstów podczas inferencji. Pomimo tych postępów, proces inferencji nadal wymaga znacznych zasobów obliczeniowych i pamięci, szczególnie dla utrzymania pamięci podręcznej Key-Value (KV), co czyni ją zasobochłonnym przedsięwzięciem. Ostatnie wysiłki badawcze koncentrowały się na poprawie wyrazistości Modeli Przestrzeni Stanu poprzez włączenie mechanizmów bramkowych zależnych od wejścia, podobnych do macierzy QKV znalezionych w mechanizmach uwagi.

Te wysiłki mają na celu zachowanie wewnętrznej liniowej postępu rekursji przestrzeni stanu, umożliwiając wydajne wykonanie za pomocą konwolucji lub wybiórczego skanowania. Ten podejście znacznie zmniejsza różnicę wydajności z transformatorami w praktycznych aplikacjach. Wśród tych postępów, Mamba wyróżnia się jako model przestrzeni stanu, który odzwierciedla cele poprzednich badań, pokazując imponujące poziomy wydajności porównywalne z transformatorami w skalach do 2,8 miliarda parametrów. Osiąga to, stosując bramkowanie zależne od wejścia do wejść rekursji modelu przestrzeni stanu (SSM), jednocześnie zapewniając wydajne obliczenia za pomocą niestandardowych jąder skanowania wybiórczego.

Modele Mieszanki Ekspertów

Modele Mieszanki Ekspertów (MoE) osiągają separację między kosztem inferencji a całkowitą liczbą parametrów, aktywując parametry w sposób wybiórczy podczas przodu. Zamiast używać wszystkich parametrów, te modele kierują tokeny do określonych ekspertów wielowarstwowego perceptronu. Idealnie, każdy ekspert jest dostosowany do przetwarzania określonego rodzaju wejścia, z mechanizmem routingu, który jest w zasadzie kompaktową siecią neuronową, decydującą o najbardziej odpowiednim ekspercie dla każdego tokenu. Ten podejście ma na celu zachowanie całościowej wyrazistości modelu o równoważnej liczbie parametrów w gęstszej konfiguracji, ale z znacznie zmniejszonymi wymogami obliczeniowymi. Zazwyczaj, router jest mapowaniem warstw liniowych z tokenów na indeksy ekspertów, przy czym każdy ekspert jest standardowym wielowarstwowym perceptronem transformatora. Niemniej jednak, deweloperzy jeszcze nie opracowali optymalnej metody treningu routera, ponieważ problem przypisania eksperta jest nieodróżnialny, a Modele Mieszanki Ekspertów często mają trudności z równowagą obciążenia i stabilnością treningu między różnymi ekspertami dla wydajności sprzętowej.

Architektura

W swojej istocie, BlackMamba wykorzystuje standardowy model transformatora składający się z bloków wielowarstwowego perceptronu i bloków uwagi dodanych sekwencyjnie wzdłuż strumienia resztowego. Teraz, większość Modeli Mieszanki Ekspertów po prostu zastępuje bloki wielowarstwowego perceptronu warstwą eksperta z routowanym. Z drugiej strony, ramy BlackMamba nie tylko zastępują blok wielowarstwowego perceptronu w transformatorze warstwą eksperta z routowanym, ale również zastępują warstwę uwagi warstwą Modelu Przestrzeni Stanu Mamba. Architektura ramy BlackMamba jest pokazana na poniższym rysunku.

Trening i Zestaw Danych

Model BlackMamba jest trenowany na ponad 300 miliardach tokenów w niestandardowym zestawie danych i wykorzystuje funkcję aktywacji SwiGLU dla wielowarstwowych perceptronów ekspertów. Ramy trenują z 8 ekspertami, liczbą, którą deweloperzy uznali za odpowiednią równowagę między śladem pamięci a kosztem inferencji modelu. Niestandardowy zestaw danych wykorzystany do treningu ramy BlackMamba składa się z mieszanki istniejących otwartych zestawów danych, w tym Starcoder, SlimPajama, Pile i innych. Poniższa tabela pokazuje wagi każdego zestawu danych wykorzystanych do treningu ramy BlackMamba. Ogółem, w zestawie danych jest 1,8 biliona tokenów.

BlackMamba: Wyniki

Aby zapewnić uczciwe porównanie między Mamba a BlackMamba, deweloperzy przeszkolili oba modele z tymi samymi parametrami treningu na tym samym danych treningowych. Ramy BlackMamba są w stanie przewyższyć zarówno Mamba, jak i modele transformatorowe w przypadku identycznego rozmiaru modelu podczas inferencji, a także w przypadku treningu operacji zmiennoprzecinkowych na sekundę. Poniższy rysunek pokazuje czas potrzebny do wygenerowania sekwencji o danej długości autoregresyjnie z jednego tokenu początkowego w funkcji długości sekwencji.

Ponadto, korzyści z opóźnienia zarówno Modelu Mieszanki Ekspertów, jak i Mamba są połączone w ramach BlackMamba, wynikając w znacznie szybszych czasach inferencji w porównaniu z modelami transformatorowymi, czystymi modelami Mamba i modelami MoE. Ponadto, zaleta inferencji ramy BlackMamba jest bezpośrednio proporcjonalna do długości sekwencji, co sprawia, że BlackMamba jest wyjątkowo skuteczny w generowaniu długich sekwencji. Przechodząc dalej, poniższy rysunek ilustruje liczbę tokenów przypisanych do modeli BlackMamba z 340 milionami i 640 milionami parametrów odpowiednio. Jak widać, większość warstw wykazuje wysoki poziom równowagi ekspertów w wyniku ulepszonych algorytmów Sinkhorn zaimplementowanych przez modele BlackMamba.

Poniższa tabela obejmuje wyniki oceny ramy BlackMamba w porównaniu z zakresem otwartych, wstępnie wytrenowanych modeli językowych. Jak można zobaczyć, ramy BlackMamba są w stanie konkurować i przewyższyć większość ram w przypadku wszystkich punktów odniesienia. Ponadto, warto zauważyć, że modele, które przewyższają BlackMamba, mają znacznie więcej parametrów, a różnica w wydajności jest minimalna, co wskazuje na zdolność ramy BlackMamba do osiągania lepszych wyników z mniejszą liczbą parametrów.

Końcowe Myśli

W tym artykule omówiliśmy BlackMamba, nową architekturę, która łączy Model Przestrzeni Stanu Mamba z modelami MoE, aby wykorzystać korzyści oferowane przez obie ramy. Eksperymenty na BlackMamba wykazały jego zdolność do przewyższenia istniejącej ramy Mamba i modeli transformatorowych w treningu FLOPs i inferencji. Wyjątkowa wydajność ramy BlackMamba pokazuje, że może ona skutecznie łączyć zdolności ram Mamba i MoE, oferując szybką i efektywną inferencję z MoE oraz generację o liniowej złożoności z Mamba. Omówiliśmy, jak architektura ramy BlackMamba jest w stanie przewyższyć silnie wytrenowane duże modele językowe, istniejącą ramę Mamba i modele MoE pod względem treningu FLOPs i kosztu inferencji. Ponadto, ramy BlackMamba również odziedziczyły generację FLOPs i zmniejszony trening z obu modeli MoE i Mamba jednocześnie.