SWARR: Pojačano učenje čini klizni prozor pažnje konkurentnim u matematičkom zaključivanju

Standardni mehanizam samopažnje (engl. self-attention) koji pokreće moderne velike jezične modele pati od temeljnog ograničenja: kvadratično skalira s duljinom konteksta. Svaki token mora se usporediti sa svakim drugim tokenom, što za duge ulaze postaje računalno preskupo. Alternativa postoji — klizni prozor pažnje (engl. sliding-window attention, SWA) koji ograničava svaki token na lokalni prozor susjednih tokena i nudi linearnu složenost — ali uz povijesno značajan pad performansi na zahtjevnim zadacima.

Osam istraživača — Kai Liu, Peijie Dong, Xinchen Xie, Jianfei Gao, Qipeng Guo, Xiaowen Chu, Shaoting Zhang i Kai Chen — objavilo je na arXiv platformi rad koji predstavlja SWARR (Sliding-Window Attention via Reinforcement Learning): dvostupanjski pristup koji taj pad performansi značajno reducira.

Problem kvadratičnog skaliranja i dugi kontekst

Puna samopažnja (engl. standard attention, SA) daje svakom tokenu pristup svim ostalim tokenima u kontekstu — što je izuzetno moćno za zadatke koji zahtijevaju praćenje dugorasponskih ovisnosti (engl. long-range dependencies). Matematičko zaključivanje upravo je takav zadatak: model mora pratiti višekoračne lance dedukcije, zadržavati privremene rezultate i vraćati se na ranije korake.

Klizni prozor pažnje to fizički ograničava: token vidi samo susjedne tokene unutar prozora fiksne veličine. Prethodni pokušaji pretvorbe modela s pune pažnje na SWA rezultirali su značajnim padom točnosti upravo na ovakvim zadacima, što je ograničavalo praktičnu primjenu ovog inače privlačnog rješenja.

Kako SWARR rješava problem u dva koraka?

SWARR je dvostupanjski pristup:

Korak 1 — Nadzirano fino podešavanje (engl. Supervised Fine-Tuning, SFT): Prethodno istrenirani model s punom samopažnjom pretvara se u model s kliznim prozorom pažnje putem standardnog SFT-a. Ovaj korak je neophodan polazišni temelj, ali autori pokazuju da sam po sebi nije dovoljan za zahtjevne zadatke zaključivanja.

Korak 2 — Adaptacija politike pojačanim učenjem (engl. Reinforcement Learning, RL): Model prolazi kroz RL fazu u kojoj se optimiziraju trajektorije zaključivanja (engl. reasoning trajectories) specifično prilagođene arhitekturalnim ograničenjima kliznog prozora. Ovo je ključna inovacija SWARR-a.

Ključni uvid: neusklađenost podataka i arhitekture

Autori identificiraju temeljni problem koji SFT ne može riješiti: neusklađenost podataka i arhitekture (engl. data-architecture mismatch). Podatkovni skupovi za treniranje standardno nastaju uz modele s punom samopažnjom, koji slobodno koriste dugorasponske ovisnosti. Kad se takav model pretvori na klizni prozor, nastaje fundamentalni raskorak — naučene strategije zaključivanja pretpostavljaju pristup informacijama koje klizni prozor fizički ne može dohvatiti.

SFT optimizira model na gotovim podacima koji su nastali bez ograničenja SWA. Pojačano učenje pak generira vlastite trajektorije zaključivanja — tzv. on-policy pristup koji je inherentno usklađen s arhitekturom. Model uči zaključivati na način koji odgovara onome što zapravo može vidjeti. Ovaj uvid autori postavljaju kao središnji doprinos rada.

Rezultati na matematičkom zaključivanju

Eksperimenti na zadacima matematičkog zaključivanja pokazuju da SWARR značajno sužava jaz između performansi modela s kliznim prozorom pažnje i onih s punom samopažnjom — uz istovremeno zadržavanje efikasnostnih dobitaka linearne složenosti pri dugom kontekstu.

Autori zaključuju da RL „temeljno mijenja zaključke o primjenjivosti kliznog prozora pažnje” na matematičke zadatke — transformirajući ga iz tek „prihvatljive” zamjene u genuino konkurentnu alternativu.

Implikacije za efikasne jezične modele

SWARR otvara put prema modelima koji mogu procesirati duge kontekste bez kvadratičnog računalnog troška, a pritom zadržati visoke performanse na zahtjevnim zadacima zaključivanja. To je relevantno za širok spektar primjena: od analize dugih pravnih dokumenata i izvornog koda, do rješavanja složenih matematičkih problema.

Šire gledano, rad pokazuje da pažljivo dizajnirani RL može biti ključni alat za prilagodbu LLM arhitektura koje su izvorno trenirane s jednim mehanizmom pažnje na arhitekturalno drugačije, efikasnije alternative — bez nužnosti ponovnog treniranja od nule.