ho trascorso alcune ore a esaminare il repository /karpathy/autoresearch riga per riga. l'angolo degli "agenti AI che fanno ricerca" è quello che sta attirando tutta l'attenzione, ma penso che la cosa più interessante sia ciò che c'è effettivamente all'interno dello script di addestramento e le decisioni ingegneristiche che rendono il ciclo di ricerca efficiente. è uno dei setup di addestramento a file singolo più densi che abbia mai letto. iniziamo con la cosa che rende possibile l'intero progetto: il budget di tempo è fissato a 300 secondi di orologio. non passi fissi, non token fissi, non flops fissi. secondi di orologio. questo sembra un dettaglio minore, ma è l'intera ragione per cui il ciclo autonomo funziona. l'agente può rendere il modello 3 volte più grande, ridurre la dimensione del batch della metà, sostituire un'architettura completamente diversa, e il risultato è ancora direttamente comparabile a ogni altro esperimento perché tutti hanno avuto esattamente 5 minuti di addestramento sulla stessa gpu. se avessi fissato i passi, un modello più grande riceverebbe meno aggiornamenti del gradiente al secondo e verrebbe penalizzato ingiustamente. se avessi fissato i token, avresti lo stesso problema. fissare il tempo di muro significa che stai ponendo la domanda giusta: dato questo hardware e questo tempo, qual è il miglior modello che puoi produrre? tutto il resto è una variabile libera. l'agente può esplorare l'intera superficie di pareto della dimensione del modello rispetto al throughput rispetto alla velocità di convergenza senza che nessuno di questi compromessi venga confuso dal protocollo di valutazione. la metrica è anche scelta con attenzione. sono bit per byte, non perdita di entropia incrociata. l'entropia incrociata dipende dalla dimensione del tuo vocabolario. un modello con 32k token e un modello con 8k token avranno valori di perdita molto diversi anche se comprimono i dati altrettanto bene. bpb normalizza questo sommando l'entropia incrociata per token in nats, sommando le lunghezze in byte utf-8 dei token target e convertendo nats-per-byte in bits-per-byte. quindi anche se l'agente cambia qualcosa che influisce sulla distribuzione effettiva dei token, il confronto rimane equo. queste due scelte, tempo di muro fisso e una metrica invariata rispetto al vocabolario, trasformano quella che sarebbe una ricerca disordinata e incomparabile in un problema di ottimizzazione pulito. ora il modello stesso. è un GPT ma con un sacco di trucchi moderni che vale la pena comprendere. prima di tutto, RMSnorm ovunque. sugli input del blocco (pre-norm), e anche su query e chiavi proprio prima del prodotto scalare di attenzione. questa cosa del QK-norm è importante perché senza di essa le norme di q e k possono crescere illimitatamente durante l'addestramento, causando un'attenuazione degli logit di attenzione e saturando il softmax. normalizzare q e k mantiene i prodotti scalari in un intervallo stabile indipendentemente da quanto sia profonda la rete o da come evolvono le dinamiche di addestramento. l'attenzione stessa è FA 3, caricata attraverso la libreria kernels. utilizza l'implementazione di varunneal su hopper (sm_90) e torna a una build della comunità su gpu più vecchie. il pattern di attenzione è "SSSL" che significa tre strati di attenzione a finestra scorrevole (finestra = metà della lunghezza della sequenza) seguiti da uno strato di attenzione causale completa, ripetendo. questo è il pattern sparso-a-denso che vedi in mistral e gemma2. gli strati di attenzione locale sono computazionalmente economici perché la matrice di attenzione è a bande, e lo strato globale periodico consente il flusso di informazioni attraverso l'intero contesto. con 8 strati e un pattern di 4 caratteri ottieni gli strati 0,1,2 locali, lo strato 3 globale, gli strati 4,5,6 locali, lo strato 7 globale. l'ultimo strato è forzato globale indipendentemente dal pattern. la cosa dell'embedding dei valori è sottile e penso sia sottovalutata. ogni altro strato ottiene la propria tabella di embedding, completamente separata dall'embedding principale dei token, che mappa gli id dei token direttamente a vettori di dimensione valore. questi vengono mescolati nei valori di attenzione attraverso un gate appreso: v = v + 2 * sigmoid(W_gate @ x:32) * ve. il peso del gate è inizializzato a zero, quindi sigmoid(0) = 0.5, moltiplicato per 2 dà 1.0, che è un punto di partenza neutro. durante l'addestramento il modello può imparare ad amplificare o sopprimere l'embedding dei valori per testa in base ai primi 32 dimensioni dello stato nascosto. questo proviene dalla linea di lavoro ResFormer e l'intuizione è che fornisce all'attenzione un collegamento diretto all'identità del token. i vettori di valore possono trasportare informazioni su "quale token è in questa posizione" senza che tali informazioni debbano sopravvivere alle trasformazioni del flusso residuo dai livelli precedenti. è essenzialmente una connessione di salto dall'input direttamente nei valori di attenzione, controllata in modo che il modello possa decidere quando è utile. ci sono anche scalari apprendibili per strato sul flusso residuo: x = lambda_residi * x + lambda_x0i * x0, dove x0 è l'embedding normalizzato dallo strato 0. ogni strato può controllare indipendentemente quanto ascolta il residuo in esecuzione rispetto all'input originale. i lambda residui partono da 1.0, i lambda x0 partono da 0.1. questa è una versione morbida dell'idea del "residuo disaccoppiato". in un trasformatore standard il flusso residuo è una somma di tutte le uscite degli strati precedenti e diventa sempre più inquinato man mano che si scende. dare a ogni strato accesso all'embedding originale pulito significa che non deve imparare a "annullare" gli strati precedenti per recuperare informazioni a basso livello. gli logit sono softcapped a 15 tramite tanh(logits/15)*15, il che impedisce al modello di essere eccessivamente sicuro all'inizio dell'addestramento quando le rappresentazioni sono ancora rumorose. ma onestamente la parte più interessante dell'intero file è l'ottimizzatore. MuonAdamW è un ottimizzatore combinato che distribuisce diverse regole di aggiornamento in base al gruppo di parametri. gli embedding (embedding dei token, embedding dei valori, testa di unembedding) e gli scalari per strato ottengono AdamW standard con tassi di apprendimento diversi per ciascun gruppo. la diffusione è selvaggia. il lr dell'embedding è 0.6, il lr dell'unembedding è 0.004, c'è una differenza di 150x, ed è intenzionale. la matrice di embedding vede ogni singolo token e deve aggiornarsi in modo aggressivo. la matrice di unembedding è una sonda lineare sulla rappresentazione finale e beneficia della stabilità. i tassi di apprendimento per embedding, embedding dei valori e unembedding sono tutti scalati da (d_model / 768)^(-0.5), che è una correzione ispirata a muP. man mano che la larghezza del modello cambia, quei tassi di apprendimento si adattano per mantenere le dinamiche di apprendimento delle caratteristiche invarianti rispetto alla scala. i tassi di apprendimento scalari per i lambda per strato vengono gestiti separatamente e non ricevono questa scalatura. gi matrice di pesi 2D nel trasformatore, proiezioni di attenzione e pesi mlp, ottengono Muon, e qui diventa genuinamente interessante. muon prende il gradiente, applica il momento di nesterov, quindi esegue un'iterazione di newton-schulz per approssimare la decomposizione polare della matrice del gradiente. la decomposizione polare fattorizza una matrice G in G = U * S dove U è ortogonale e S è simmetrica positiva semi-definita. muon calcola U, la matrice ortogonale più vicina al gradiente, e la usa come direzione di aggiornamento. l'iterazione di newton-schulz è di 5 passi. per matrici alte (più righe che colonne), A = X^T @ X quindi X -> aX + X @ (bA + cA^2). per matrici larghe, A = X @ X^T quindi X -> aX + (bA + cA^2) @ X. i coefficienti sono hardcoded da una precomputation. lo chiamano "polar express." l'intero processo si compila in un singolo kernel fuso tramite torch.compile. perché questo è importante? perché per le matrici di pesi il gradiente della norma di frobenius (quello che adam e sgd usano) è geometricamente sbagliato. la "corretta" direzione di discesa più ripida per una matrice di pesi è quella che minimizza la perdita soggetta al vincolo che l'aggiornamento abbia norma spettrale unitaria, non norma di frobenius unitaria. il fattore polare ortogonale ti dà esattamente questo. in pratica significa che muon fa aggiornamenti effettivi molto più grandi perché non sta sprecando la dimensione del passo sulla scalatura dei valori singolari. ruota solo quelli. questo è il motivo per cui muon converge significativamente più velocemente di adam sulle matrici di pesi del trasformatore. muon mantiene anche buffer di momento per elemento (stessa forma dei parametri, impilati attraverso ciascun gruppo di forma), ma a differenza di adam non tiene traccia dei secondi momenti per elemento. le stime del secondo momento sono per riga o per colonna dopo l'ortogonalizzazione, non per elemento. è qui che entra in gioco NorMuon. sopra la base muon c'è NorMuon, uno schema di riduzione della varianza. dopo l'ortogonalizzazione, calcola le stime del secondo momento per riga (o per colonna a seconda del rapporto di aspetto), mantiene una media mobile esponenziale di quelle e ridimensiona l'aggiornamento in modo che ogni dimensione di output ottenga la propria dimensione di passo adattiva. è essenzialmente l'idea di adattività di adam ma applicata nel sistema di coordinate ortogonalizzato piuttosto che nello spazio dei parametri grezzi. il decadimento del peso è anche non standard. è "cauteloso", il che significa che deca solo i parametri dove la direzione di aggiornamento di muon concorda con il segno del parametro: maschera = (g * params) >= 0. questo evita il noto modo di fallimento in cui il decadimento del peso spinge i parametri verso zero contro i desideri dell'aggiornamento, il che può destabilizzare l'addestramento. un piccolo dettaglio che ho apprezzato: dopo il primo passo di addestramento, il codice chiama gc.collect(), gc.freeze(), gc.disable() per spegnere completamente il garbage collector di python. il GC di python gira periodicamente e causa pause di ~500ms. quando il tuo budget totale è di 300 secondi e ogni passo è forse di 300ms, una pausa casuale del GC ti costa quasi 2 passi di addestramento. attivano manualmente gc.collect() ogni 5000 passi come compromesso. questo è il tipo di cosa che impari solo profilando le vere esecuzioni di addestramento e notando misteriosi cali di throughput. i primi 11 passi (da 0 a 10) non vengono conteggiati nemmeno nel budget di tempo. questo è il riscaldamento in cui torch.compile fa il suo lavoro e i kernel CUDA vengono JIT'd. senza questa esclusione, esperimenti diversi avrebbero ricevuto quantità diverse di "vero" addestramento a seconda di quanto tempo impiega la compilazione per quella particolare configurazione del modello. ancora una volta, una scelta di design che sembra piccola ma è critica per rendere comparabili gli esperimenti. ora zoom out. il vero ciclo di autoresearch è: l'agente legge program.md (un file markdown che descrive il suo lavoro), modifica train.py, fa commit, esegue per 5 minuti, controlla se val_bpb è migliorato, mantiene o ripristina, ripete. program.md dice esplicitamente "NON FERMARE MAI." l'agente funziona indefinitamente fino a quando l'umano non lo uccide. ~12 esperimenti all'ora, ~100 durante la notte mentre dormi. ...