Astra i looped transformers. Jak model ponownie używa warstw

W doniesieniach o Astrze pojawiło się określenie looped transformer. Sebastian Raschka wyjaśnia ten mechanizm na przykładzie otwartego modelu Nanbeige4.2-3B.

Te same wagi, kolejny przebieg

Nanbeige przepuszcza dane dwukrotnie przez stos 22 warstw. Drugi przebieg przetwarza wynik pierwszego, korzystając z tych samych wag, czyli parametrów wyuczonych podczas treningu. Model wykonuje więc 44 operacje warstw bez przechowywania drugiej kopii parametrów.

Taki układ wymaga mniej więcej dwukrotnie większej pracy samego stosu niż pojedynczy przebieg. Nie jest jednak równoważny 44 niezależnym warstwom: współdzielone wagi ograniczają swobodę uczenia różnych przekształceń. Faktyczny czas odpowiedzi zależy także od wykonania operacji i obsługi pamięci.

Autorzy Nanbeige wybrali dwa przebiegi jako najlepszy kompromis w swoich eksperymentach. Więcej powtórzeń dawało niewielką poprawę, spowalniało trening i pogarszało jego stabilność. Lepiej sprawdzało się też uczenie z pętlą od początku niż dodawanie jej do wcześniej wytrenowanego modelu.

Poza wagami model przechowuje pamięć podręczną KV: klucze i wartości dla wcześniejszych tokenów, używane przez mechanizm uwagi. Nanbeige zachowuje osobne wpisy dla obu przebiegów. Ich współdzielenie zmniejszało pamięć KV o połowę, lecz dawało gorsze wyniki.

Liczbę powtórzeń można dobierać do tokena

Universal Transformer z 2018 roku korzystał ze współdzielonych parametrów i pozwalał kończyć obliczenia osobno dla poszczególnych pozycji w sekwencji.

Mixture-of-Recursions dodaje uczony router, który przydziela tokenom liczbę przejść przez wspólny blok. Część kończy przetwarzanie wcześniej, inne otrzymują dodatkowe obliczenia. Autorzy łączą to z selektywnym przechowywaniem wpisów KV lub ich współdzieleniem. Badanie obejmuje uczenie na FineWeb-Edu i skale modeli bazowych od 135 mln do 1,7 mld parametrów; dostosowanie routera do rozbudowanych łańcuchów rozumowania pozostaje dalszą pracą.

Co z widocznością rozumowania

W chain of thought model generuje pośrednie tokeny tekstu. Pętla dodaje obliczenia w stanach ukrytych przed wyemitowaniem kolejnego tokena. Oba mechanizmy mogą współistnieć: Nanbeige obsługuje tryby Think i Non-Think, opisane w raporcie treningowym. Raschka sugeruje, że dodatkowe obliczenia wewnętrzne mogą ograniczyć potrzebę generowania pośrednich tokenów, ale sama pętla tego nie gwarantuje.

W karcie systemowej Astry OpenAI raportuje spadek skuteczności nadzoru opartego na tekście rozumowania. Osobno opisuje wzrost kontroli modelu nad własnym CoT podczas uczenia ze wzmocnieniem; tej konkretnej zmiany nie przypisuje zmianom architektury. Karta nie ujawnia układu warstw Astry, więc przykład Nanbeige wyjaśnia możliwy mechanizm, lecz nie potwierdza konstrukcji modelu OpenAI.

Wykorzystuję treści generowane przez AI jako część mojego codziennego procesu nauki.