Я провів кілька годин, переглядаючи /karpathy/autoresearch, рядок за рядком. Тема «агенти ШІ, які проводять дослідження», привертає всю увагу, але, на мою думку, цікавіше — це те, що насправді всередині навчального скрипту та інженерні рішення, які роблять пошук напруженим. Це одна з найщільніших однорядних навчальних систем, які я читав. Дозвольте почати з того, що робить увесь проєкт можливим: часовий бюджет встановлений на 300 секундах на стіні. Не фіксовані кроки, не фіксовані токени, не фіксовані флопси. Настінний годинник секунди. Це звучить як дрібниця, але саме через це працює автономна петля. Агент може зробити модель у 3 рази більшою, скоротити розмір партії вдвічі, замінити зовсім іншу архітектуру, і результат все одно напряму порівнянний з усіма іншими експериментами, бо всі вони отримали рівно 5 хвилин навчання на одній відеокарті. Якби ви виправили кроки, більша модель отримувала б менше оновлень градієнтів за секунду, і ви б несправедливо її карали. Якби ви виправили токени, виникла б та сама проблема. Час ремонту стіни означає, що ви ставите правильне питання: враховуючи це обладнання і стільки часу, яку модель найкраще ви можете виготовити? все інше — це вільна змінна. Агент може досліджувати повну парето-поверхню розміру моделі, пропускної здатності та швидкості збіжності, не збиваючи з пантелику жоден із цих компромісів протоколом оцінки. метрика також ретельно підбирається. Це біти на байт, а не втрати через перехресну ентропію. Крос-ентропія залежить від розміру вашого словникового запасу. Модель з 32k токенів і модель з 8k токенів матимуть дуже різні значення втрат, навіть якщо вони однаково добре стискають дані. BPB нормалізує це, підсумовуючи крос-ентропію на один токен у NAT, сумуючи довжини байтів UTF-8 цільових токенів і конвертуючи NAT-на байт у біти на байт. Тож навіть якщо агент змінює щось, що впливає на ефективний розподіл токенів, порівняння залишається справедливим. Ці два варіанти — фіксований час на стіні та незмінна до словникової метрики — перетворюють те, що могло б бути заплутаним непорівнянним пошуком, на чисту задачу оптимізації. Тепер про саму модель. це GPT, але з купою сучасних хитрощів, які варто зрозуміти. По-перше, RMSnorm всюди. на блочних входах (до норми), а також на запитах і ключах безпосередньо перед добутком крапки уваги. Ця норма QK-норма важлива, бо без неї норми Q і K можуть безмежно зростати під час тренувань, що призводить до загострення логів уваги та насичення softmax. Нормалізація Q і K підтримує добутки DOT у стабільному діапазоні незалежно від глибини мережі чи розвитку динаміки тренувань. сама увага — це FA 3, завантажена через бібліотеку ядра. Він використовує реалізацію Varunneal на Hopper (sm_90) і повертається до спільноти, побудованої на старіших GPU. патерн уваги — «SSSL», що означає три шари уваги ковзного вікна (вікно = половина довжини послідовності), за якими слідує один шар повної причинної уваги, що повторюється. Це той самий візерунок, який ви бачите в Містралі та Гемма2. Локальні шари уваги обчислювально дешеві, оскільки матриця уваги є смугоподібною, а періодичний глобальний шар дозволяє інформації циркулювати через повний контекст. З 8 шарами та 4-символьним візерунком ви отримуєте шари 0,1,2 локальні, шари 3 глобальні, шари 4,5,6 локальні, шари 7 глобальні. Останній шар примусово глобальний незалежно від патерну. Ідея з вкладенням цінності тонка і, на мою думку, недооцінена. Кожен інший шар має власну таблицю вкладення, повністю відокремлену від основного вкладення токена, яка безпосередньо відображає ідентифікатори токена у вектори значення-розмірності. Вони змішуються з значеннями уваги через вивчений гейт: v = v + 2 * сигмоїд(W_gate @ x:32) * ve. Вага затвора ініціалізована нуль, тому sigmoid(0) = 0,5, помножено на 2 дає 1,0, що є нейтральною відправною точкою. Під час навчання модель може навчитися підсилювати або пригнічувати вкладення значень на голову на основі перших 32 вимірів прихованого стану. це з напрямку роботи ResFormer, і інтуїція полягає в тому, що він дає увагу прямий короткий шлях до ідентичності токенів. Вектори значень можуть нести інформацію про «який токен знаходиться в цій позиції», і ця інформація не повинна переживати залишкові потокові трансформації з попередніх шарів. По суті, це пропускний зв'язок від входу безпосередньо до значень уваги, обмежений, щоб модель могла вирішувати, коли це корисно. Існують також щошарові вивчені скаляри на залишковому потоку: x = lambda_residi * x + lambda_x0i * x0, де x0 — нормалізоване вкладення з шару 0. Кожен шар може незалежно контролювати, наскільки він слухає залишковий сигнал порівняно з початковим входом. Залишкові лямбди починаються з 1,0, X0 — з 0,1. Це м'яка версія ідеї «розплутаного залишку». У стандартному трансформаторі залишковий потік — це сума всіх попередніх виходів шарів, і він все більше забруднюється глибше. Надання кожному шару доступу до чистого оригінального вбудовування означає, що йому не потрібно вчитися «скасовувати» попередні шари, щоб відновити низькорівневу інформацію. Логіти мають софткап на 15 через Tanh(Logits/15)*15, що запобігає надмірній самовпевненості моделі на ранніх етапах навчання, коли представлення ще гучні. Але чесно кажучи, найцікавіша частина всього файлу — це оптимізатор. MuonAdamW — це комбінований оптимізатор, який розподіляє різні правила оновлення залежно від групи параметрів. Embeddings (embedding token, value embedding, unembedding head) та скаляри на рівні отримують стандартний AdamW з різними швидкостями навчання для кожної групи. Розповсюдження просто шалене. Embedding LR — 0.6, unembedding LR — 0.004, це різниця у 150 разів, і це зроблено навмисно. Матриця вбудовування бачить кожен токен і потребує агресивного оновлення. Матриця без вкладення є лінійним зонд на фінальному представленні і має перевагу стабільності. швидкість навчання вбудовування, вкладення значень і невкладення масштабуються на (d_model / 768)^(-0.5), що є корекцією, натхненною muP. Зі зміною ширини моделі ці швидкості навчання коригуються так, щоб динаміка навчання ознак залишалася інваріантною до масштабу. Скалярні швидкості навчання для лямбда на шар обробляються окремо і не мають такого масштабування. 2D-матриці ваги в трансформері, проекції уваги та ваги MLP — отримують Muon, і тут починається справді цікаво. Мюон приймає градієнт, застосовує імпульс Нестерова, а потім виконує ітерацію Ньютона-Шульца для наближення полярного розкладу градієнтної матриці. полярний розклад розкладає матрицю G на G = U * S, де U — ортогональна, а S — симетрична додатна напіввизначена. мюон обчислює U, найближчу ортогональну матрицю до градієнта, і використовує її як напрямок оновлення. Ітерація Ньютона-Шульца складається з 5 кроків. для високих матриць (більше рядків, ніж стовпців) A = X^T @ X, потім X -> aX + X @ (bA + cA^2). для широких матриць A = X @ X^T, тоді X -> aX + (bA + cA^2) @ X. Коефіцієнти жорстко закодовані з попереднього обчислення. Вони називають це «Polar Express». Все компілюється в одне злите ядро через torch.compile. Чому це важливо? Бо для вагових матриць градієнт норми Фробеніуса (який використовують Адам і SGD) геометрично неправильний. «Правильний» найкрутіший напрямок спуску для матриці ваг — це той, що мінімізує втрати за умови обмеження, що оновлення має одиничну спектральну норму, а не одиничну норму Фробеніуса. Ортогональний полярний коефіцієнт дає саме це. На практиці це означає, що мюон робить набагато більші ефективні оновлення, бо не витрачає розмір кроку на масштабування одиничних значень. Він лише обертає їх. Ось чому мюон збігається значно швидше, ніж у Адама на матрицях ваги трансформатора. Мюон зберігає буфери імпульсу на елемент (ту ж форму, що й параметри, накладені по кожній групі фігур), але, на відміну від Адама, він не відстежує секундні моменти для кожного елемента. Оцінки другого моменту відбуваються за рядок або стовпець після ортогоналізації, а не за елементами. ось тут і з'являється NorMuon. поверх базового мюону є NorMuon — схема зменшення дисперсії. Після ортогоналізації він обчислює оцінки секундного моменту за рядком (або стовпцем залежно від співвідношення сторін), підтримує експоненціальне ковзне середнє їх і масштабує оновлення, щоб кожна вихідна вимірність отримала власний адаптивний розмір кроку. Це, по суті, ідея адаптивності Адама, але застосована в ортогоналізованій системі координат, а не у сирому просторі параметрів. Згасання ваги також є нестандартним. Він «обережний», тобто зменшує лише параметри, де напрямок оновлення мюонів збігається зі знаком параметра: Mask = (g * параметри) >= 0. Це дозволяє уникнути відомого режиму відмови, коли зменшення ваги змушує параметри наближатися до нуля, всупереч бажанням оновлення, що може дестабілізувати тренування. Одна маленька деталь, яку я оцінив: після першого етапу навчання код викликає gc.collect(), gc.freeze(), gc.disable(), щоб повністю вимкнути сміттєвий збирач Python. GC Python працює періодично і викликає ~500 мс застій. коли ваш загальний бюджет — 300 секунд, а кожен крок — близько 300 мс, випадкова пауза GC коштує майже 2 тренувальних кроків. Вони вручну запускають gc.collect() кожні 5000 кроків як компроміс. Це те, що можна дізнатися лише, профілюючи реальні тренувальні заїзди та помічаючи загадкові втрати пропускної здатності. Перші 11 кроків (від 0 до 10) також не враховуються у часовому бюджеті. Це розігрів, коли torch.compile робить свою справу, а ядра CUDA піддаються JIT. Без цього виключення різні експерименти отримували б різний рівень «реального» навчання залежно від того, скільки часу займає компіляція для конкретної конфігурації моделі. Знову ж таки, це дизайнерський вибір, який здається невеликим, але критично важливий для порівняння експериментів. Тепер віддаляйте масштаб. Фактичний цикл автодослідження такий: агент читає program.md (файл markdown, що описує його завдання), змінює .py навчання, фіксує, працює 5 хвилин, перевіряє, чи val_bpb покращено, зберігає або повертає, повторює. program.md чітко каже: «НІКОЛИ НЕ ЗУПИНЯЙСЯ». Агент біжить безкінечно, поки людина не вбиває його. ~12 експериментів на годину, ~100 за ніч, поки спите. ...