Strávil jsem pár hodin procházením /karpathy/autoresearch repozitáře řádek po řádku. Téma "AI agentů dělajících výzkum" je to, co přitahuje veškerou pozornost, ale myslím, že zajímavější je to, co je skutečně uvnitř tréninkového skriptu a inženýrská rozhodnutí, která dělají vyhledávací smyčku těsnou. Je to jeden z nejhustších tréninkových systémů s jedním souborem, které jsem četl. Začnu tím, co celý projekt umožňuje: časový rozpočet je pevně stanoven na 300 sekund na nástěnné hodiny. Ne pevné kroky, ne pevné žetony, ne pevné flopy. Nástěnné hodiny sekundy. Zní to jako drobnost, ale právě proto autonomní smyčka funguje. Agent může model zvětšit třikrát, snížit velikost šarže na polovinu, přidat úplně jinou architekturu a výsledek je stále přímo srovnatelný s ostatními experimenty, protože všechny měly přesně 5 minut tréninku na stejné GPU. Kdybyste místo toho nastavili kroky, větší model by měl méně aktualizací gradientu za sekundu a nespravedlivě byste ho penalizovali. Kdybyste opravili žetony, měli byste stejný problém. Oprava času na zdi znamená, že kladete správnou otázku: s tímto hardwarem a tolika časem, jaký je nejlepší model, který můžete vyrobit? Všechno ostatní je volná proměnná. Agent může prozkoumat celou Pareto plochu velikosti modelu vs propustnost vs rychlost konvergence, aniž by tyto kompromisy narušil evaluační protokol. metrika je také pečlivě zvolena. Je to bity na bajt, ne ztráta křížové entropie. Křížová entropie závisí na velikosti slovní zásoby. Model s 32k tokeny a model s 8k tokeny budou mít velmi odlišné hodnoty ztrát, i když data komprimují stejně dobře. BPB to normalizuje sečtením křížové entropie na token v NATS, sečtem délk UTF-8 bajtů cílových tokenů a převodem NAT-per-byte na bity na bajt. Takže i když agent změní něco, co ovlivňuje efektivní rozdělení tokenů, srovnání zůstává spravedlivé. Tyto dvě možnosti, pevný čas na zdi a metrika invariantní slovní zásoby, proměňují to, co by bylo chaotické, nesrovnatelné hledání, v čistý optimalizační problém. Teď samotný model. Je to GPT, ale s řadou moderních triků, které stojí za pochopení. za prvé, RMSnorm všude. na vstupech bloků (před normou), a také na dotazech a klíčích těsně před součinem bodů pozornosti. Tato QK-norma je důležitá, protože bez ní se normy Q a K mohou během tréninku nekonečně zvyšovat, což způsobuje zostření pozornostních logitů a nasycení softmaxu. Normalizace Q a K udržuje skalární součiny v stabilním rozsahu bez ohledu na hloubku sítě nebo jak se vyvíjejí dynamika tréninku. samotná pozornost je FA 3, načtená přes knihovnu Kernels. Používá implementaci Varunnealu na Hopperu (sm_90) a vrací se k komunitě postavené na starších GPU. vzor pozornosti je "SSSL", což znamená tři vrstvy posuvné okénkové pozornosti (okno = polovina délky sekvence), následované jednou vrstvou plné kauzální pozornosti, která se opakuje. To je vzor řídké až husté, který vidíte v Mistralu a Gemma2. Lokální vrstvy pozornosti jsou výpočetně levné, protože matice pozornosti je pásmovaná a periodická globální vrstva umožňuje tok informací napříč celým kontextem. S 8 vrstvami a 4znakovým vzorem dostanete vrstvy 0,1,2 lokální, vrstvu 3 globální, vrstvy 4,5,6 lokální a vrstvu 7 globální. poslední vrstva je nuceně globální bez ohledu na vzor. To vkládání hodnoty je jemné a myslím, že je nedoceněné. Každá druhá vrstva má vlastní tabulku vkládání, zcela oddělenou od hlavního vnoření tokenů, která mapuje ID tokenů přímo na vektory s hodnotovou dimenzí. Tyto hodnoty se míchají do hodnot pozornosti přes naučenou bránu: V = V + 2 * Sigmoid(W_gate @ x:32) * ve. Hmotnost hradel je nulově inicializovaná, takže sigmoid(0) = 0,5, krát 2 dává 1,0, což je neutrální výchozí bod. Přetrénování se model může naučit zesílit nebo potlačit hodnotu vložení na hlavu na základě prvních 32 dimenzí skrytého stavu. to pochází z oblasti práce ResFormer a intuice je, že dává pozornost přímé zkratce k identitě tokenu. Vektory hodnoty mohou přenášet informace o tom, "jaký token je na této pozici", aniž by tato informace musela přežít zbytkové transformace proudu z předchozích vrstev. V podstatě jde o přeskočení spojení ze vstupu přímo do hodnot pozornosti, které je nastaveno jako gate, aby model mohl rozhodnout, kdy je to užitečné. Na zbytkovém proudu jsou také naučitelné skaláry na vrstvě: x = lambda_residi * x + lambda_x0i * x0, kde x0 je normalizované vložení z vrstvy 0. Každá vrstva může nezávisle ovládat, jak moc naslouchá běžícímu reziduu oproti původnímu vstupu. reziduální lambdy začínají na 1,0, x0 lambdy začínají na 0,1. To je měkká verze myšlenky "rozpleteného rezidua". Ve standardním transformátoru je zbytkový proud součtem všech předchozích výstupů vrstev a čím hlouběji se posouváte, tím více znečišťuje. Poskytnutí přístupu každé vrstvě k čistému původnímu embeddingu znamená, že se nemusí učit "zrušovat" předchozí vrstvy, aby získala nízkoúrovňové informace. Logity jsou softcapovány na 15 pomocí TANH(logits/15)*15, což zabraňuje tomu, aby byl model příliš sebevědomý na začátku trénování, když jsou reprezentace stále hlučné. Ale upřímně, nejzajímavější částí celého souboru je optimalizátor. MuonAdamW je kombinovaný optimalizátor, který vysílá různá pravidla aktualizace na základě skupiny parametrů. embeddingy (token embedding, value embeddings, unembedding head) a skaláry na úrovni per level mají standardní AdamW s různými učebními rychlostmi pro každou skupinu. Spread je divoký. Vložení LR je 0,6, odembedding LR 0,004, to je rozdíl 150x a je to záměrné. Embedding matice vidí každý jeden token a musí se agresivně aktualizovat. Matice bez vložení je lineární sonda na konečné reprezentaci a těží ze stability. Rychlosti učení vkládání, vkládání hodnoty a odembedding učení jsou škálovány pomocí (d_model / 768)^(-0,5), což je korekce inspirovaná muP. Jak se mění šířka modelu, tyto rychlosti učení se přizpůsobují tak, aby dynamika učení prvků zůstala škálově invariantní. Skalární rychlosti učení pro lambdy na jednotlivých vrstvách jsou zpracovávány samostatně a toto škálování nedosahují. 2D váhové matice v transformeru, projekce pozornosti a váhy MLP dostanou Muon, a tady to začíná být opravdu zajímavé. Muon vezme gradient, aplikuje Nesterovovu hybnost a poté provede Newton-Schulzovu iteraci k aproximaci polárního rozkladu gradientní matice. polární rozklad rozkladá matici G na G = U * S, kde U je ortogonální a S je symetricky kladně semidefinitní. muon vypočítá U, nejbližší ortogonální matici ke gradientu, a použije ji jako směr aktualizace. Newton-Schulzova verze má 5 kroků. pro vysoké matice (více řádků než sloupců) platí A = X^T @ X, pak X -> aX + X @ (bA + cA^2). pro široké matice platí A = X @ X^T pak X -> aX + (bA + cA^2) @ X. koeficienty jsou pevně zakódovány z předpočtu. Říkají tomu "Polar Express." Celý systém se kompiluje do jednoho fúzovaného jádra přes Torch.compile. Proč na tom záleží? Protože u váhových matic je Frobeniův normální gradient (který používají Adam a SGD) geometricky nesprávný. "Správný" nejstrmější směr sestupu pro hmotnostní matici je ten, který minimalizuje ztrátu za předpokladu, že aktualizace má jednotkovou spektrální normu, nikoli jednotkovou Frobeniovu normu. Ortogonální polární faktor vám přesně toto umožňuje. V praxi to znamená, že Muon provádí mnohem větší efektivní aktualizace, protože neztrácí velikost kroků škálováním singulárních hodnot. jen je otáčí. Proto se muon konverguje výrazně rychleji než Adam na maticích hmotnosti transformátorů. Muon sice udržuje momentální buffery pro jednotlivé prvky (stejný tvar jako parametry, naskládané do každé skupiny tvarů), ale na rozdíl od Adamu nesleduje momenty pro jednotlivé prvky v sekundách. Odhady druhého momentu jsou po ortogonalizaci na řádek nebo sloupec, nikoli na prvek. a právě zde přichází na řadu NorMuon. nad základním mionem je NorMuon, schéma redukce rozptylu. Po ortogonalizaci vypočítá odhady druhého momentu na řádek (nebo sloupec podle poměru stran), udržuje exponenciální klouzavý průměr těchto a aktualizuje aktualizaci tak, aby každá výstupní dimenze dostala vlastní adaptivní velikost kroku. V podstatě jde o Adamovu adaptivitu, ale aplikovanou v ortogonalizovaném souřadnicovém systému, nikoli v surovém parametrovém prostoru. Pokles hmotnosti je také nestandardní. Je to "opatrné", což znamená, že parametry klesají pouze tam, kde směr aktualizace mionu odpovídá znaménku parametru: mask = (g * parametry) >= 0. Tím se vyhne známému režimu selhání, kdy pokles hmotnosti tlačí parametry k nule proti přání aktualizace, což může destabilizovat trénink. Jeden malý detail, který jsem ocenil: po úplně prvním tréninkovém kroku kód zavolá gc.collect(), gc.freeze(), gc.disable() a úplně vypne pythonův sběrač odpadu. Pythonův GC běží periodicky a způsobuje ~500 ms zastavení. když je váš celkový rozpočet 300 sekund a každý krok třeba 300 ms, náhodná pauza GC vás stojí téměř 2 tréninkové kroky. Ručně spouštějí GC.collect() každých 5000 kroků jako kompromis. Tohle je něco, co se naučíte jen profilováním skutečných tréninkových běhů a zaznamenáním záhadných poklesů průtoku. Prvních 11 kroků (0 až 10) se také nezapočítává do časového rozpočtu. to je ten rozběh, kdy Torch.compile dělá svou práci a CUDA jádra jsou JIT-ována. Bez této výjimky by různé experimenty dostávaly různé množství "skutečného" tréninku v závislosti na tom, jak dlouho trvá kompilace pro danou konfiguraci modelu. Opět designová volba, která se zdá malá, ale je zásadní pro srovnání experimentů. Teď oddál pohled. Skutečná smyčka autovýzkumu je: Agent čte program.md (Markdown soubor, který popisuje jeho úkol), upraví train .py, commituje, běží 5 minut, kontroluje, zda val_bpb zlepšeno, ponechá nebo vrátí zpět, opakuje. program.md výslovně říká "NIKDY NEPŘESTÁVAT." agent běží neomezeně dlouho, dokud ho člověk nezabije. ~12 experimentů za hodinu, ~100 přes noc, když spíš. ...