YaRN: jak rozciągnąć kontekst, którego model nie widział

Przewijasz kartę modelu na Hugging Face i w tabelce mruga do ciebie „128k context”. Sto dwadzieścia osiem tysięcy tokenów — cała książka wleci w jeden prompt. Serce rośnie. A potem, gdzieś na dole, drobnym druczkiem: model bazowy trenowany na sekwencjach 4096. No i człowiek głupieje. Jak to? Uczyłeś go patrzeć na cztery tysiące tokenów naraz, a teraz ma ogarnąć trzydzieści dwa razy więcej? Skąd on niby wie, jak wygląda token numer 100 000, skoro w życiu takiego nie widział?

To nie jest pytanie z gatunku czepialstwa. Sam pewnie nieraz widziałeś, jak model z „długim kontekstem” pięknie streszcza pierwsze pięć stron, a od dziesiątej zaczyna zmyślać, mieszać imiona i gubić wątek — dokładnie w miejscu, w którym teoretycznie miał błyszczeć. Bo rozciąganie okna kontekstu to nie jest przesunięcie suwaka. To operacja na sercu tego, jak model w ogóle rozumie „gdzie” jest dany token. I tu wchodzi YaRN — metoda z pracy grupy z Nous Research, która robi to tanio i, co ważniejsze, robi to dobrze.¹

Skąd model w ogóle wie, który token jest który

Zacznijmy od tego, że transformer sam z siebie nie ma pojęcia o kolejności. Wrzuć mu „pies goni kota” i „kota goni pies” — bez informacji o pozycji to dla niego ta sama torba słów. Trzeba mu jakoś powiedzieć, że ten token jest pierwszy, tamten setny.

Współczesne modele — LLaMA, Mistral, Qwen — robią to przez RoPE (rotary position embeddings, rotacyjne osadzenia pozycyjne) z pracy Su i innych.² Pomysł jest elegancki: zamiast doklejać pozycję jako osobny wektor, RoPE obraca wektory zapytań i kluczy o kąt proporcjonalny do pozycji tokena. Każdy wymiar embeddingu kręci się z inną częstotliwością — jedne szybko, jak wskazówka sekundowa, inne wolno, jak godzinowa. Token na pozycji 5 jest obrócony trochę, token na pozycji 5000 — dużo. Model uczy się czytać te kąty i z różnicy obrotów między dwoma tokenami wnioskuje, jak daleko od siebie leżą.

I teraz clou problemu: model widział te wskazówki tykające tylko do pozycji 4096. Dalej jest terra incognita. Kąty, które w treningu nigdy się nie pojawiły. Puść go na token 50 000, a wskazówka sekundowa zakręci się tyle razy, że model nie ma bladego pojęcia, co z tym zrobić — to jak pokazać komuś zegar, który obrócił się poza tarczę, w miejsce, którego nigdy nie oznaczono.

Dlaczego „po prostu ściśnij pozycje” psuje robotę

Pierwszy odruch jest oczywisty. Skoro model zna zakres do 4096, a chcemy 64k, to przeskalujmy pozycje liniowo — token 64 000 udawaj, że jesteś tokenem 4000. Ściśnij całą oś czasu tak, żeby zmieściła się w znanym zakresie. To jest Position Interpolation (PI, interpolacja pozycji) z pracy Chena i innych — i na papierze brzmi rozsądnie.

Problem w tym, że PI traktuje wszystkie częstotliwości RoPE jednakowo — zwalnia i sekundnik, i wskazówkę godzinową o ten sam czynnik.³ A to katastrofa dla tych najszybszych wymiarów. To właśnie one odpowiadają za rozróżnianie tokenów leżących tuż obok siebie — „czy słowo A jest bezpośrednio przed B, czy jest między nimi jeszcze jedno”. Ściśnij je razem z resztą, a model traci rozdzielczość na najdrobniejszym poziomie. Efekt? Jakość spada nawet na krótkim kontekście, który przecież działał bez zarzutu. Rozciągnąłeś okno i po drodze rozmyłeś to, co model umiał od początku. Klasyczny handel, w którym oddajesz więcej, niż dostajesz.

Jak YaRN dzieli częstotliwości na pasma

No to jak zrobić to mądrzej? Skoro problem jest w tym, że traktujemy wszystkie częstotliwości jednakowo — to przestańmy. To jest cała intuicja YaRN.

Zanim doszli do wersji finalnej, po drodze była NTK-aware interpolation (interpolacja świadoma NTK) — zamiast liniowo ściskać pozycje, zmienia się podstawę RoPE z b na b * s^(|D|/(|D|-2)), przez co „nacisk” interpolacji rozkłada się nierówno po wymiarach.³ Lepiej, ale wciąż z grubsza jednym pociągnięciem po wszystkim.

YaRN idzie krok dalej z podejściem NTK-by-parts (interpolacja pasmowa) — i to jest serce metody. Dzielisz wymiary częstotliwości na pasma za pomocą funkcji rampy γ(r):³

• wysokie częstotliwości (te lokalne, szybkie wskazówki — sąsiad przy sąsiedzie) zostawiasz w spokoju, bez interpolacji. Ekstrapolujesz. Niech tykają jak tykały.
• niskie częstotliwości (te globalne, wolne — gruby zarys „to jest gdzieś na końcu dokumentu”) interpolujesz normalnie, bo to one muszą pomieścić nowy, większy zakres.
• pasmo pośrodku rampa miesza płynnie, żeby nie było skoku na styku.

Sedno: nie ruszasz tego, co model umie najlepiej (drobna rozdzielczość lokalna), a rozciągasz tylko to, co i tak opisuje wielką skalę. Zegar dostaje większą tarczę, ale sekundnik dalej odmierza sekundy tak samo dokładnie.

Drugi bezpiecznik: temperatura uwagi

I tu jest smaczek, który łatwo przegapić, bo nie ma nic wspólnego ze skalowaniem częstotliwości. YaRN dorzuca attention temperature scaling — skalowanie temperatury w softmaxie uwagi.³ Kiedy rozciągasz kontekst, rozkład uwagi robi się „płaski”, model rozprasza się na tysiące tokenów zamiast celować. Więc mnożysz logity uwagi przez stały współczynnik, wyliczony z prostego, empirycznego wzoru:

sqrt(1/t) = 0.1 * ln(s) + 1

gdzie s to współczynnik skalowania kontekstu (dla 4096 → 64k mamy s = 16). Piękne w tym jest to, że działa jak zwykły mnożnik na logitach — nie dotyka wag modelu, nie wymaga dodatkowego treningu, wchodzi za darmo. Osobne pokrętło, doklejone z boku, a robi sporo dobrego dla ostrości. Nie pomyl go ze skalowaniem samego RoPE — to dwa różne mechanizmy, które YaRN łączy w jedno.

Tanio, czyli ile to naprawdę kosztuje

Teraz najlepsze — bo tu YaRN naprawdę błyszczy. To nie jest trening od zera. To nawet nie jest porządny fine-tuning na górze danych. To krótkie douczenie już gotowego modelu, żeby oswoił się z nowymi kątami.

Ile krótkie? W pracy LLaMA-2 7B i 13B douczono jakieś 400 kroków przy s = 16 (4096 → 64k), plus dodatkowe ~200 kroków przy s = 32, żeby dobić do 128k — razem około 600 kroków, na danych PG19 w segmentach po 64k, przy batchu 64.³ Objętościowo to szacunkowo ~0,1% oryginalnego korpusu pretreningowego. Jedna dziesiąta procenta. W liczbach z abstraktu: YaRN potrzebuje około 10x mniej tokenów douczenia i 2,5x mniej kroków treningowych niż wcześniejsze metody.¹

I jeszcze jedna rzecz, która robi wrażenie: modele ekstrapolują poza długość, na której je douczono. Serię Mistral 7B rozciągnięto z 8k do 64k i 128k tą samą metodą — a trening na segmentach 64k dawał poprawne działanie aż do 128k.³ Model, który w douczeniu nie widział nic dłuższego niż 64k, radzi sobie dwa razy dalej. To już nie jest tylko „nauczyliśmy go nowego zakresu” — to „nauczyliśmy go reguły, którą sam rozciąga dalej”.

Zrób to sam: YaRN w llama.cpp

I teraz uwaga na rozróżnienie, na którym łatwo się wyłożyć. Wszystko powyżej — kroki douczenia, dane PG19 — to warstwa treningu opisana w papierze. Ale jest druga warstwa: inferencja. llama.cpp implementuje skalowanie YaRN jako parametry runtime, które przekręcasz przy uruchomieniu, bez żadnego douczania:⁴

llama-server -m model-yarn-128k.gguf \
  -c 65536 \
  --yarn-orig-ctx 8192 \
  --yarn-ext-factor 1.0 \
  -fa on

• --yarn-orig-ctx — oryginalny kontekst treningowy modelu (domyślnie 0, czyli „wczytaj z metadanych GGUF”). To punkt odniesienia, względem którego liczy się s.
• --yarn-ext-factor — współczynnik mieszania ekstrapolacji (domyślnie -1.00; 0.0 to pełna interpolacja).

Jeśli model został wypuszczony jako wariant YaRN, metadane zwykle same podpowiedzą sensowne wartości — wtedy w ogóle nie musisz w to grzebać. A jak flagi w llama.cpp w ogóle działają i które z tej setki naprawdę musisz znać, rozbieraliśmy w osobnym wpisie. Jest jeszcze wariant Dynamic NTK — aktualizuje współczynnik skalowania s = max(1, l'/L) w locie, per krok inferencji, zamiast trzymać stałą wartość ustaloną raz.³ Sprytne, bo krótkie prompty nie płacą podatku za rozciągnięcie, którego akurat nie potrzebują.

Cała ta rodzina metod, swoją drogą, ładnie zazębia się z tym, jak w ogóle liczy się uwaga — a to temat, który drążyliśmy przy FlashAttention.

Większa tarcza, ten sam sekundnik

Wróćmy na koniec do tego zegara. Naiwna interpolacja bierze tarczę i ściska ją tak, że sekundnik zlewa się z minutową — niby zmieściłeś więcej godzin, ale nie odczytasz już, czy jest 12:01, czy 12:02. YaRN robi odwrotnie: zostawia sekundnik w spokoju, a rozciąga tylko wolne wskazówki, które i tak odmierzają grube kawałki. Dokładasz kilkaset kroków douczenia — tyle, ile trzeba, żeby model oswoił nowe kąty — i nagle patrzy na sto dwadzieścia osiem tysięcy tokenów tak, jakby zawsze umiał. Nie dlatego, że zobaczył każdą nową pozycję. Dlatego, że ktoś mądrze przekręcił mu zegar, nie tłukąc szkła.