RWKV-X: Nowa architektura dla długich kontekstów

Wciąż rosnące zapotrzebowanie na przetwarzanie dłuższych i bardziej złożonych sekwencji przesunęło granice Large Language Models (LLMs). Tradycyjne architektury oparte na Transformerach, choć potężne, zmagają się z poważnymi problemami skalowalności ze względu na ich kwadratową złożoność w odniesieniu do długości sekwencji. To ograniczenie staje się szczególnie widoczne przy obsłudze rozszerzonych danych wejściowych kontekstu, utrudniając ich zdolność do efektywnego przechwytywania i wykorzystywania informacji z odległych części sekwencji. W odpowiedzi na to wyzwanie pojawiła się fala innowacyjnych podejść, mających na celu osiągnięcie liniowej złożoności w przetwarzaniu długich sekwencji.

Metody te obejmują modele Linear Attention, State Space Models (takie jak Mamba), Linear RNNs (jak DeltaNet) i RWKV. Każda z tych architektur oferuje unikalne rozwiązanie problemu kwadratowej złożoności, umożliwiając bardziej efektywne przetwarzanie długich sekwencji. Jednak te liniowe architektury często napotykają trudności w pełnym zrozumieniu i wykorzystaniu informacji o długim kontekście.

Na przykład RWKV-7 (model z 2,9 miliardami parametrów) wykazuje wysoką dokładność w zadaniach wyszukiwania passkey do 28 tys. tokenów. Jednak jego wydajność gwałtownie spada powyżej tego progu. Nawet przy ciągłym wstępnym uczeniu przy użyciu danych o długości 128 tys., ograniczenia dotyczące długiego kontekstu utrzymują się. Ten problem nie jest unikalny dla RWKV; rozciąga się na inne architektury, takie jak Mamba, stanowiąc fundamentalne wyzwanie dla tej klasy modeli. Walka o utrzymanie wydajności w rozszerzonych kontekstach podkreśla kluczowy obszar do poprawy w modelach językowych o liniowej złożoności.

Krajobraz modeli językowych o liniowej złożoności

Modele językowe o liniowej złożoności pojawiły się jako kuszące alternatywy dla architektur opartych na transformatorach, omijając kwadratowe obciążenia obliczeniowe związane z przetwarzaniem długich sekwencji. Rodzina modeli RWKV, wyróżniająca się w tej dziedzinie, mistrzowsko łączy paralelizm transformatorów podczas treningu z reprezentacją stanu rekurencyjnego podobnego do RNN.

Ewolucja RWKV obejmuje kilka iteracji, począwszy od fundamentalnego RWKV-4, przechodząc do RWKV-5, RWKV-6 i kulminując w RWKV-7. Każda iteracja przyniosła udoskonalenia i ulepszenia, zwiększając możliwości modelu i eliminując ograniczenia. Ponadto, hybrydowe modele językowe, takie jak Jamba, Zamba i MiniMax, odcisnęły swoje piętno, wprowadzając unikalne hybrydowe projekty, jeszcze bardziej wzbogacając krajobraz modeli o liniowej złożoności.

Dążenie do wydajnego przetwarzania długiego kontekstu doprowadziło również do rozwoju innowacyjnych mechanizmów uwagi. Native Sparse Attention, na przykład, organizuje tokeny w bloki czasowe, wykorzystując trzy odrębne ścieżki uwagi: skompresowane, zgrubne tokeny dla globalnego kontekstu, selektywnie zachowane, szczegółowe tokeny dla lokalnych szczegółów i przesuwane okna do przechwytywania lokalnych informacji kontekstowych. Inne godne uwagi mechanizmy uwagi obejmują SeerAttention i Block Attention (MoBA), każdy oferujący unikalne strategie zwracania uwagi na istotne informacje w długich sekwencjach.

RWKV-X: Hybrydowa architektura dla ulepszonego modelowania kontekstu dalekiego zasięgu

Naukowcy z Guangdong Laboratory of Artificial Intelligence and Digital Economy (SZ), Shenzhen, Hohai University, Nanjing, Shenzhen University i Qinghai University, Xining, wprowadzili nową hybrydową architekturę o nazwie RWKV-X. Architektura ta pomysłowo łączy wydajność RWKV w modelowaniu zależności krótkiego zasięgu ze sparse attention mechanism specjalnie zaprojektowanym do przechwytywania długiego kontekstu.

W przeciwieństwie do poprzednich podejść hybrydowych, RWKV-X osiąga złożoność liniową w czasie podczas treningu i złożoność stałą w czasie podczas dekodowania wnioskowania. To czyni go wyjątkowo wydajnym w przetwarzaniu długich sekwencji. Model wykazuje niemal idealną dokładność na teście porównawczym wyszukiwania passkey 64K, gdy jest wstępnie trenowany na sekwencjach 64K tokenów w sposób ciągły. Konsekwentnie przewyższa poprzednie modele RWKV-7 w testach porównawczych długiego kontekstu, zachowując jednocześnie wysoką wydajność w zadaniach krótkiego kontekstu.

Innowacje w RWKV-X stanowią znaczący krok naprzód w rozwiązywaniu wyzwań związanych z modelowaniem języka długiego kontekstu. Łącząc mocne strony modeli rekurencyjnych i mechanizmów sparse attention, RWKV-X osiąga równowagę między wydajnością a dokładnością, torując drogę do bardziej efektywnego przetwarzania rozszerzonych sekwencji.

RWKV-X: Architektura i trening

RWKV-X ucieleśnia hybrydową architekturę, integrując bloki RWKV-7 z blokami sparse attention, aby wykorzystać mocne strony obu podejść. Zamiast trenować od zera, RWKV-X opiera się na istniejących modelach, wykorzystując przeplataną metodę rozszerzania bloków i mechanizm zerowej inicjalizacji inspirowany LLaMA Pro.

Proces treningowy składa się z dwóch etapów, starannie zaprojektowanych w celu optymalizacji wydajności modelu zarówno w krótkich, jak i długich kontekstach:

  • Wstępny trening krótkiego kontekstu: Początkowo model jest trenowany na krótkich kontekstach 1024 tokenów wyekstrahowanych z zestawu danych MiniPile. Podczas tego etapu wszystkie parametry z wyjątkiem tych w nowo dodanych blokach są zamrożone, co zapewnia zachowanie wstępnie wytrenowanej wiedzy z podstawowego modelu RWKV-7. Pozwala to nowo dodanym blokom na dostosowanie się do istniejącej architektury bez zakłócania wstępnie wytrenowanych reprezentacji.
  • Ciągły trening długiego kontekstu: Drugi etap obejmuje ciągły trening długiego kontekstu przy użyciu zestawu danych ProLong-64K i długości kontekstu 64K tokenów, przetwarzając łącznie około 1 miliarda tokenów. Podczas tej fazy wszystkie parametry są odmrażane i wspólnie optymalizowane, co pozwala modelowi na dostrojenie swoich reprezentacji i uczenie się zależności dalekiego zasięgu. Trening wykorzystuje stratę Long-context Cross-Entropy (LongCE), która dynamicznie waży tokeny w oparciu o ich ważność. Ta funkcja straty pomaga modelowi skupić się na najbardziej istotnych częściach sekwencji, poprawiając jego zdolność do przechwytywania relacji dalekiego zasięgu.

Dwuetapowy proces treningowy pozwala RWKV-X na efektywne połączenie wydajności RWKV-7 dla modelowania krótkiego zasięgu ze świadomością kontekstu dalekiego zasięgu mechanizmu sparse attention. Poprzez wstępne trenowanie w krótkich kontekstach, a następnie dostrajanie w długich kontekstach, model uczy się efektywnie integrować informacje z różnych części sekwencji.

RWKV-X: Ocena i wydajność

Ocena krótkiego kontekstu ujawnia, że RWKV-X utrzymuje konkurencyjną wydajność w standardowych testach porównawczych, demonstrując swoją zdolność do efektywnego radzenia sobie z krótszymi sekwencjami. Mniejszy RWKV-X (0,22B) osiąga średni wynik 51,0, porównywalny z 51,8 RWKV-7. W większej skali RWKV-X (3,6B) osiąga 71,9, ściśle dopasowując się do RWKV-7 (2,9B, 72,8) i Qwen2.5-3B (71,4), jednocześnie przewyższając LLaMA3.2-3B (69,7). Wyniki te potwierdzają skuteczność RWKV-X jako uniwersalnego szkieletu LLM bez poświęcania wydajności w krótszych kontekstach.

Ponadto analiza wydajności demonstruje doskonałe charakterystyki skalowania RWKV-X dla długich sekwencji. Przy 128 tys. tokenów RWKV-X osiąga przyspieszenie 1,37 raza w porównaniu z Flash-Attention v3, przy czym ta przewaga rozszerza się wraz ze wzrostem długości kontekstu. Wskazuje to, że RWKV-X staje się coraz bardziej wydajny w porównaniu z innymi mechanizmami uwagi wraz ze wzrostem długości sekwencji.

Wysoka wydajność RWKV-X zarówno w krótkich, jak i długich kontekstach podkreśla jego wszechstronność i wydajność jako modelu językowego. Jego zdolność do utrzymywania konkurencyjnej wydajności w krótszych sekwencjach przy jednoczesnym osiąganiu znacznych przyspieszeń w dłuższych sekwencjach czyni go obiecującą architekturą dla szerokiego zakresu zastosowań.

RWKV-X: Ograniczenia i przyszłe kierunki

RWKV-X jawi się jako hybrydowy model językowy, który z powodzeniem łączy wydajność RWKV w modelowaniu zależności krótkiego zasięgu z nowym mechanizmem sparse attention zaprojektowanym specjalnie do modelowania kontekstu dalekiego zasięgu. Chociaż RWKV-X demonstruje wysoką wydajność i efektywność w modelowaniu języka długiego kontekstu, pozostaje kilka ograniczeń.

Po pierwsze, jego mechanizm sparse attention, który opiera się na selekcji chunków top-k, wykorzystuje podejście heurystyczne, które może przeoczyć semantycznie istotne zależności. Strategia selekcji top-k może nie zawsze wychwytywać najważniejsze informacje w sekwencji, potencjalnie prowadząc do suboptymalnej wydajności.

Po drugie, obecna implementacja pokazuje, że dekodowanie sparse attention działa wolniej niż vanilla RWKV, co wskazuje, że potrzebne są dalsze wysiłki inżynieryjne w celu optymalizacji wydajności. Chociaż RWKV-X osiąga znaczne przyspieszenia w porównaniu z innymi mechanizmami uwagi w długich sekwencjach, jego dekodowanie sparse attention jest nadal wolniejsze niż vanilla RWKV, co sugeruje, że istnieje pole do poprawy w jego implementacji.

Przyszłe badania mogłyby skupić się na rozwiązaniu tych ograniczeń poprzez eksplorację bardziej wyrafinowanych mechanizmów sparse attention, optymalizację implementacji dekodowania sparse attention i badanie alternatywnych strategii treningowych. Pokonując te wyzwania, RWKV-X ma potencjał, aby stać się jeszcze potężniejszym i wydajniejszym modelem językowym dla zastosowań w długim kontekście.