Mamba-3: Ewolucja architektury SSM rzuca wyzwanie transformatorom
W sektorze rozwoju dużych modeli językowych (LLM) środek ciężkości przesuwa się w stronę optymalizacji czasu inferencji. Choć architektury oparte na transformatorach pozostają rynkowym standardem, ich kwadratowa złożoność obliczeniowa staje się coraz wyraźniejszym wąskim gardłem. Odpowiedzią na te wyzwania jest Mamba-3 – najnowsze osiągnięcie badaczy z Carnegie Mellon University, Princeton, Together AI oraz Cartesia AI. Projekt ten redefiniuje fundamenty modeli stanów przestrzennych (SSM), stawiając na podejście typu „inference-first”.
Precyzja drugiego rzędu: dyskretyzacja trapezoidalna
Modele SSM pracują w systemie czasu ciągłego, co w praktyce oznacza konieczność ich dyskretyzacji do przetwarzania ciągów danych. Dotychczasowe wersje architektury Mamba polegały na heurystyce wykładniczego Eulera – przybliżeniu pierwszego rzędu. Mamba-3 wprowadza jednak dyskretyzację wykładniczo-trapezoidalną. Ta zmiana matematyczna pozwala na uzyskanie dokładności drugiego rzędu przy aproksymacji całki stanu wejściowego.
Co to oznacza w praktyce? Rekurencja modelu ewoluowała z układu dwuelementowego do trójelementowego. Efekt ten odpowiada zastosowaniu niejawnego splotu o szerokości 2 bezpośrednio wewnątrz rdzenia obliczeniowego. Dzięki temu Mamba-3 eliminuje potrzebę stosowania zewnętrznych, krótkich splotów przyczynowych, które w poprzednich generacjach były niezbędne do stabilnej pracy modeli rekurencyjnych.
Przełamywanie barier logicznych dzięki liczbom zespolonym
Jedną z pięt achillesowych modeli opartych na liczbach rzeczywistych była niezdolność do rozwiązywania specyficznych zadań logicznych, takich jak śledzenie parzystości bitów. Wynika to z ograniczeń macierzy przejścia – liczby rzeczywiste nie są w stanie efektywnie reprezentować dynamiki rotacyjnej, kluczowej dla operacji modulo czy arytmetyki bitowej. Mamba-3 wprowadza zespolone modele SSM, które teoretycznie zrównują się z modelami rzeczywistymi wykorzystującymi adaptacyjne rotacyjne osadzenia pozycyjne (RoPE).
Zastosowanie tak zwanego „triku RoPE” pozwala modelowi na agregowanie rotacji zależnych od danych w czasie. W testach empirycznych Mamba-3 bezbłędnie radzi sobie z zadaniami syntetycznymi, podczas gdy jej poprzednicy oraz inne warianty rzeczywiste wykazywały skuteczność na poziomie losowego zgadywania.
Wydajność sprzętowa: architektura MIMO
Tradycyjne podejście Single-Input Single-Output (SISO) w modelach SSM cierpiało na niską intensywność obliczeniową – nowoczesne procesory graficzne, takie jak H100, „nudziły się” podczas dekodowania ze względu na bariery transferu pamięci. Mamba-3 przechodzi na strukturę Multi-Input Multi-Output (MIMO), zwiększając rangę projekcji wejściowych i wyjściowych.
Przejście na MIMO przekształca aktualizację stanu z prostego iloczynu wektorowego w pełnoprawne mnożenie macierzy. Pozwala to na czterokrotne zwiększenie liczby operacji zmiennoprzecinkowych (FLOPs) przy zachowaniu tego samego rozmiaru stanu. Kluczowe jest to, że dodatkowe obliczenia odbywają się „w cieniu” trwającego przesyłu danych z pamięci I/O, co przekłada się na wyższą jakość modelowania bez odczuwalnego wzrostu latencji.
Wyniki i perspektywy
Analiza wyników na zbiorze FineWeb-Edu potwierdza skuteczność nowej architektury. Wersja Mamba-3 o rozmiarze 1,5 miliarda parametrów przewyższa modele Mamba-2 oraz Gated DeltaNet. Co więcej, Mamba-3 osiąga podobną gęstość (perplexity) jak poprzednia generacja, wykorzystując przy tym o połowę mniejszy stan. To istotny krok w kierunku oszczędnego zarządzania zasobami podczas pracy z długim kontekstem.
Dzięki zastosowaniu optymalizowanych kerneli Triton i CuTe, Mamba-3 udowadnia, że głębokie zmiany w fundamentach matematycznych modeli są kluczem do pogodzenia teoretycznej efektywności subkwadratowej z praktycznymi wymaganiami nowoczesnych aplikacji AI.
