spędziłem kilka godzin przeglądając repozytorium /karpathy/autoresearch linia po linii. a „agenci AI prowadzący badania” to temat, który przyciąga całą uwagę, ale myślę, że ciekawszą rzeczą jest to, co tak naprawdę znajduje się w skrypcie treningowym oraz decyzje inżynieryjne, które sprawiają, że pętla wyszukiwania jest zwarta. To jeden z najbardziej gęstych jednolitych zestawów treningowych, jakie czytałem. Zacznę od rzeczy, która sprawia, że cały projekt jest możliwy: budżet czasowy jest ustalony na 300 sekund zegarowych. Nie ustalone kroki, nie ustalone tokeny, nie ustalone flopy. Sekundy zegarowe. To brzmi jak drobny szczegół, ale to cała przyczyna, dla której autonomiczna pętla działa. Agent może zwiększyć model 3x, zmniejszyć rozmiar partii o połowę, wprowadzić zupełnie inną architekturę, a wynik nadal będzie bezpośrednio porównywalny z każdym innym eksperymentem, ponieważ wszystkie miały dokładnie 5 minut treningu na tym samym GPU. Gdybyś ustalił kroki, większy model otrzymałby mniej aktualizacji gradientu na sekundę i byłby niesprawiedliwie karany. Gdybyś ustalił tokeny, miałbyś ten sam problem. Ustalenie czasu ściany oznacza, że zadajesz właściwe pytanie: biorąc pod uwagę ten sprzęt i ten czas, jaki jest najlepszy model, jaki możesz wyprodukować? Wszystko inne to zmienna swobodna. Agent może badać pełną powierzchnię Pareto rozmiaru modelu w porównaniu do przepustowości i szybkości zbieżności, bez żadnych z tych kompromisów zakłócających protokół oceny. Metryka jest również starannie dobrana. To bity na bajt, a nie strata entropii krzyżowej. Entropia krzyżowa zależy od rozmiaru słownika. Model z 32k tokenów i model z 8k tokenów będą miały bardzo różne wartości straty, nawet jeśli kompresują dane równie dobrze. bpb normalizuje to, sumując entropię krzyżową na token w nats, sumując długości bajtów utf-8 docelowych tokenów i przekształcając nats na bajty na bity na bajt. Więc nawet jeśli agent zmienia coś, co wpływa na efektywny rozkład tokenów, porównanie pozostaje sprawiedliwe. Te dwa wybory, ustalony czas ściany i metryka niezależna od słownika, przekształcają to, co byłoby chaotycznym, nieporównywalnym wyszukiwaniem, w czysty problem optymalizacji. Teraz sam model. To GPT, ale z wieloma nowoczesnymi sztuczkami, które warto zrozumieć. Po pierwsze, RMSnorm wszędzie. Na wejściach bloku (pre-norm), a także na zapytaniach i kluczach tuż przed iloczynem punktowym uwagi. Ta rzecz QK-norm jest ważna, ponieważ bez niej normy q i k mogą rosnąć nieograniczenie podczas treningu, powodując, że logity uwagi stają się ostre, a softmax nasyca się. Normalizowanie q i k utrzymuje iloczyny punktowe w stabilnym zakresie, niezależnie od tego, jak głęboka jest sieć lub jak rozwijają się dynamiki treningowe. Sama uwaga to FA 3, załadowana przez bibliotekę kernels. Używa implementacji varunneala na hopperze (sm_90) i przechodzi do wersji społecznościowej na starszych GPU. Wzór uwagi to „SSSL”, co oznacza trzy warstwy uwagi z przesuwanym oknem (okno = połowa długości sekwencji), a następnie jedna warstwa pełnej uwagi przyczynowej, powtarzając. To jest wzór rzadko-gęsty, który widzisz w mistralu i gemma2. Lokalne warstwy uwagi są obliczeniowo tanie, ponieważ macierz uwagi jest pasmowa, a okresowa warstwa globalna pozwala na przepływ informacji przez pełny kontekst. Przy 8 warstwach i wzorze 4-znakowym otrzymujesz warstwy 0,1,2 lokalne, warstwę 3 globalną, warstwy 4,5,6 lokalne, warstwę 7 globalną. Ostatnia warstwa jest wymuszona globalnie, niezależnie od wzoru. Rzecz z osadzeniem wartości jest subtelna i myślę, że niedoceniana. Każda inna warstwa ma swoją własną tabelę osadzenia, całkowicie oddzieloną od głównego osadzenia tokenów, która mapuje identyfikatory tokenów bezpośrednio na wektory wymiarów wartości. Te są mieszane w wartości uwagi przez nauczoną bramkę: v = v + 2 * sigmoid(W_gate @ x:32) * ve. Waga bramki jest inicjowana na zero, więc sigmoid(0) = 0.5, razy 2 daje 1.0, co jest neutralnym punktem wyjścia. W trakcie treningu model może nauczyć się wzmacniać lub tłumić osadzenie wartości na głowę w zależności od pierwszych 32 wymiarów stanu ukrytego. To pochodzi z linii pracy ResFormer, a intuicja jest taka, że daje uwadze bezpośredni skrót do tożsamości tokenu. Wektory wartości mogą przenosić informacje o „jakim tokenie jest na tej pozycji”, bez konieczności, aby ta informacja przetrwała transformacje strumienia resztkowego z wcześniejszych warstw. To zasadniczo połączenie skip z wejścia bezpośrednio do wartości uwagi, bramkowane, aby model mógł zdecydować, kiedy jest to przydatne. Są też uczące się skalarne na każdej warstwie w strumieniu resztkowym: x = lambda_residi * x + lambda_x0i * x0, gdzie x0 to znormalizowane osadzenie z warstwy 0. Każda warstwa może niezależnie kontrolować, jak bardzo słucha bieżącego strumienia resztkowego w porównaniu do oryginalnego wejścia. Lambdy resztkowe zaczynają się od 1.0, lambdy x0 zaczynają się od 0.1. To jest miękka wersja pomysłu „rozdzielonej resztki”. W standardowym transformatorze strumień resztkowy jest sumą wszystkich wyjść wcześniejszych warstw i staje się coraz bardziej zanieczyszczony w miarę głębokości. Dając każdej warstwie dostęp do czystego oryginalnego osadzenia, nie musi uczyć się „cofać” wcześniejszych warstw, aby odzyskać informacje na niskim poziomie. Logity są miękko ograniczone do 15 przez tanh(logits/15)*15, co zapobiega nadmiernej pewności modelu na początku treningu, gdy reprezentacje są nadal hałaśliwe. Ale szczerze mówiąc, najciekawsza część całego pliku to optymalizator. MuonAdamW to połączony optymalizator, który przydziela różne zasady aktualizacji w zależności od grupy parametrów. Osadzenia (osadzenie tokenów, osadzenia wartości, głowa unembedding) i skalarne na każdej warstwie otrzymują standardowy AdamW z różnymi współczynnikami uczenia dla każdej grupy. Rozrzut jest dziki. Współczynnik uczenia osadzenia wynosi 0.6, współczynnik uczenia unembedding wynosi 0.004, to różnica 150x, i jest to zamierzone. Macierz osadzenia widzi każdy pojedynczy token i musi aktualizować agresywnie. Macierz unembedding to liniowy sondowanie na końcowej reprezentacji i korzysta ze stabilności. Współczynniki uczenia osadzenia, osadzenia wartości i unembedding są skalowane przez (d_model / 768)^(-0.5), co jest poprawką inspirowaną muP. W miarę zmiany szerokości modelu, te współczynniki uczenia dostosowują się, aby utrzymać dynamikę uczenia cech w skali. Współczynniki uczenia dla lambd na każdej warstwie są obsługiwane osobno i nie podlegają temu skalowaniu. Macierze wag 2D w transformatorze, projekcje uwagi i wagi mlp, otrzymują Muon, i tutaj robi się naprawdę interesująco. muon bierze gradient, stosuje moment Nesterova, a następnie wykonuje iterację Newtona-Schulza, aby przybliżyć dekompozycję biegunową macierzy gradientu. Dekompozycja biegunowa dzieli macierz G na G = U * S, gdzie U jest ortogonalna, a S jest symetryczna dodatnio półdefiniowana. muon oblicza U, najbliższą ortogonalną macierz do gradientu, i używa tego jako kierunku aktualizacji. Iteracja Newtona-Schulza to 5 kroków. Dla wysokich macierzy (więcej wierszy niż kolumn), A = X^T @ X, a następnie X -> aX + X @ (bA + cA^2). Dla szerokich macierzy, A = X @ X^T, a następnie X -> aX + (bA + cA^2) @ X. Współczynniki są zakodowane na stałe z wstępnych obliczeń. Nazywają to „polar express”. Całość kompiluje się do jednego zintegrowanego jądra za pomocą torch.compile. Dlaczego to ma znaczenie? Ponieważ dla macierzy wag gradient normy Frobeniusa (co używa adam i sgd) jest geometrycznie błędny. „Poprawny” kierunek najszybszego spadku dla macierzy wag to ten, który minimalizuje stratę przy założeniu, że aktualizacja ma jednostkową normę spektralną, a nie jednostkową normę Frobeniusa. Ortogonalny czynnik biegunowy daje dokładnie to. W praktyce oznacza to, że muon dokonuje znacznie większych efektywnych aktualizacji, ponieważ nie marnuje rozmiaru kroku na skalowanie wartości osobliwych. Tylko je obraca. Dlatego muon zbiega się znacznie szybciej niż adam na macierzach wag transformatora. muon utrzymuje również bufory momentów dla każdego elementu (w tej samej formie co parametry, ułożone w każdej grupie kształtów), ale w przeciwieństwie do adama nie śledzi drugich momentów dla każdego elementu. Szacowania drugiego momentu są na wierszach lub kolumnach po ortogonalizacji, a nie dla każdego elementu. To jest miejsce, w którym wchodzi NorMuon. Na bazie muon jest NorMuon, schemat redukcji wariancji. Po ortogonalizacji oblicza szacowania drugiego momentu na wierszach (lub kolumnach w zależności od proporcji) i utrzymuje wykładniczą średnią ruchomą tych, a następnie przeskalowuje aktualizację, aby każdy wymiar wyjściowy otrzymał własny adaptacyjny rozmiar kroku. To zasadniczo pomysł adaptacyjności adama, ale zastosowany w ortogonalizowanym układzie współrzędnych, a nie w surowej przestrzeni parametrów. Wygaszenie wag jest również niestandardowe. Jest „ostrożne”, co oznacza, że wygasza tylko parametry, gdzie kierunek aktualizacji muon zgadza się z znakiem parametru: maska = (g * params) >= 0. To unika znanego trybu awarii, w którym wygaszenie wag popycha parametry w kierunku zera wbrew życzeniom aktualizacji, co może destabilizować trening. Jeden mały szczegół, który doceniłem: po pierwszym kroku treningowym kod wywołuje gc.collect(), gc.freeze(), gc.disable(), aby całkowicie wyłączyć zbieracz śmieci Pythona. GC Pythona działa okresowo i powoduje opóźnienia rzędu ~500ms. Kiedy twój całkowity budżet wynosi 300 sekund, a każdy krok trwa może 300ms, losowa przerwa GC kosztuje ci prawie 2 kroki treningowe. Ręcznie wywołują gc.collect() co 5000 kroków jako kompromis. To jest coś, czego uczysz się tylko przez profilowanie rzeczywistych sesji treningowych i zauważanie tajemniczych spadków wydajności. Pierwsze 11 kroków (0 do 10) również nie jest liczonych w budżecie czasowym. To rozgrzewka, podczas której torch.compile robi swoje, a jądra CUDA są JIT'd. Bez tego wyłączenia różne eksperymenty miałyby różne ilości „prawdziwego” treningu w zależności od tego, jak długo trwa kompilacja dla danej konfiguracji modelu. Jeszcze raz, wybór projektowy, który wydaje się mały, ale jest kluczowy dla porównywalności eksperymentów. Teraz przyjrzyjmy się z szerszej perspektywy. Rzeczywista pętla autoresearch to: agent czyta program.md (plik markdown, który opisuje jego zadanie), modyfikuje train.py, zatwierdza, uruchamia przez 5 minut, sprawdza, czy val_bpb się poprawił, zachowuje lub cofa, powtarza. program.md wyraźnie mówi „NIGDY NIE ZATRZYMUJ SIĘ”. Agent działa w nieskończoność, dopóki człowiek go nie zabije. ~12 eksperymentów na godzinę, ~100 przez noc, podczas snu. ...