Naukowcy z MBZUAI i CMU prezentują Bi-Mamba: Wydajną i Skalowalną Architekturę 1-bit Mamba dla Dużych Modeli Językowych (780M, 1.3B i 2.7B Parametrów)

Postęp w dziedzinie uczenia maszynowego i nowe modele językowe

Rozwój uczenia maszynowego przyniósł ogromne postępy w tworzeniu modeli językowych, które stanowią fundament zadań takich jak generowanie tekstu czy odpowiadanie na pytania. Kluczowe znaczenie w tej dziedzinie mają architektury oparte na transformatorach oraz modele przestrzeni stanów (SSM – State-Space Models). Chociaż transformery są szeroko stosowane, ich efektywność maleje w przypadku przetwarzania długich sekwencji z powodu złożoności kwadratowej. W miarę wydłużania sekwencji zapotrzebowanie na pamięć i moc obliczeniową staje się wręcz nieosiągalne. Aby rozwiązać ten problem, badacze opracowali alternatywne podejścia, takie jak model Mamba – model przestrzeni stanów o liniowej złożoności, który zapewnia wydajność i skalowalność w zadaniach wymagających przetwarzania długich kontekstów.

Wyzwania związane z dużą skalą modeli językowych

Modele językowe o dużej skali, liczące miliardy parametrów, stawiają ogromne wyzwania związane z kosztami obliczeniowymi. Choć model Mamba oferuje korzyści wynikające z liniowej złożoności, jego rosnąca wielkość prowadzi do znacznego zużycia energii i wysokich kosztów szkolenia, co utrudnia wdrożenie. Problem ten jest szczególnie widoczny w przypadku modeli opartych na architekturach, takich jak GPT, które zazwyczaj są trenowane i wdrażane w pełnej precyzji (np. FP16 lub BF16). W odpowiedzi na rosnące zapotrzebowanie na wydajne i skalowalne systemy AI, badacze coraz częściej koncentrują się na metodach ekstremalnej kwantyzacji, które umożliwiają praktyczne zastosowanie modeli w środowiskach o ograniczonych zasobach.

Kwantyzacja i optymalizacja modeli

Jednym z podejść do redukcji kosztów obliczeniowych jest kwantyzacja, czyli zmniejszenie liczby bitów reprezentujących wagi modelu. Technika ta pozwala na kompresję modeli bez znaczącego pogorszenia ich wydajności. Większość dotychczasowych badań koncentruje się jednak na modelach opartych na transformatorach, pomijając modele przestrzeni stanów, takie jak Mamba. Zachowanie tych modeli w warunkach ekstremalnej kwantyzacji pozostaje obszarem, który wymaga dalszych badań, co tworzy lukę w opracowywaniu skalowalnych i wydajnych modeli przestrzeni stanów do zastosowań praktycznych.

Bi-Mamba – innowacyjna architektura binarna

W odpowiedzi na te wyzwania, badacze z Mohamed bin Zayed University of Artificial Intelligence oraz Carnegie Mellon University stworzyli model Bi-Mamba – skalowalną architekturę Mamba zoptymalizowaną do pracy z 1-bitową reprezentacją danych. Bi-Mamba wykorzystuje binaryzację uwzględniającą proces szkolenia (binarization-aware training), co pozwala na ekstremalną kwantyzację przy jednoczesnym zachowaniu konkurencyjnej wydajności. Model został opracowany w trzech wersjach o różnych rozmiarach: 780 milionów, 1,3 miliarda oraz 2,7 miliarda parametrów. Szkolenie przeprowadzono od podstaw z zastosowaniem autoregresyjnej straty destylacyjnej, przy wsparciu modeli nauczycieli o wysokiej precyzji, takich jak LLaMA2-7B, co zapewniło wysoką jakość wyników.

Architektura i proces szkolenia Bi-Mamba

Architektura Bi-Mamba łączy selektywną binaryzację modułów liniowych z zachowaniem pełnej precyzji w pozostałych komponentach modelu. Projekcje wejściowe i wyjściowe zostały zbinaryzowane za pomocą modułów FBI-Linear, które integrują skalowalne i uczące się parametry, umożliwiając precyzyjne odwzorowanie wag w modelu. Szkolenie odbywało się na 32 kartach graficznych NVIDIA A100 z wykorzystaniem ogromnych zbiorów danych, obejmujących 1,26 biliona tokenów pochodzących z takich źródeł jak RefinedWeb czy StarCoder.

Rezultaty eksperymentów

Przeprowadzone eksperymenty wykazały przewagę Bi-Mamba nad istniejącymi modelami. Na zestawach danych takich jak Wiki2, PTB i C4 model osiągnął wyniki odpowiednio 14,2, 34,4 i 15,0 w metryce perplexity, znacząco przewyższając alternatywy, takie jak GPTQ czy Bi-LLM, które wykazywały nawet 10-krotnie wyższe wartości perplexity. Dodatkowo, Bi-Mamba uzyskał wysokie wyniki w zadaniach zero-shot, osiągając dokładność 44,5% dla modelu 780M, 49,3% dla modelu 2,7B oraz 46,7% dla wariantu 1,3B na zbiorach takich jak BoolQ czy HellaSwag.

Kluczowe wnioski z badań

• Optymalizacja pamięci: Bi-Mamba pozwala na ponad 80% kompresji pamięci w porównaniu z modelami pełnobitowymi, zmniejszając rozmiar z 5,03 GB do 0,55 GB dla wariantu 2,7B.
• Stabilna wydajność: Mimo znacznej redukcji wymagań pamięciowych, model zachowuje porównywalną wydajność do pełnobitowych odpowiedników.
• Skalowalność: Architektura Bi-Mamba umożliwia efektywne szkolenie modeli o różnych rozmiarach, oferując konkurencyjne wyniki nawet dla największych wariantów.
• Innowacyjne podejście do binaryzacji: Dzięki selektywnej binaryzacji modułów liniowych, model unika degradacji wydajności typowej dla prostych metod binarizacji.

Podsumowanie

Bi-Mamba to przełomowe rozwiązanie, które adresuje kluczowe wyzwania związane ze skalowalnością i efektywnością modeli językowych. Dzięki zastosowaniu binaryzacji uwzględniającej proces szkolenia oraz innowacyjnym optymalizacjom architektonicznym, badacze udowodnili, że modele przestrzeni stanów mogą osiągać wysoką wydajność przy ekstremalnej kwantyzacji. Bi-Mamba nie tylko zmniejsza zużycie energii i zasobów, ale także otwiera nowe możliwości w zakresie wdrażania dużych modeli w środowiskach o ograniczonych zasobach. Wyniki badań podkreślają potencjał tego podejścia jako kroku milowego w kierunku bardziej zrównoważonych i wydajnych technologii sztucznej inteligencji.