Tilde Research 發布了 Aurora,一個新的「leverage-aware」優化器,修復了 Muon 中一個隱藏的 bug —— 該優化器一直在悄悄地為 modded-nanoGPT speedruns 和越來越多的前沿訓練管道提供動力。Bug:在像 MLP 層這樣的高矩陣中,Muon 的極因子更新創建行範數各向異性,導致一些神經元接收大規模更新,而其他神經元幾乎不接收任何更新。在第 500 個訓練步驟時,超過四分之一的神經元實際上已死。Aurora 以大約 6% 的額外計算開銷防止這種情況,作為 Muon 的替代品插入,並附帶 GitHub 上的開放程式碼。
診斷比修復更重要。NorMuon(一個先前的中間版本)透過事後歸一化到單位範數來糾正行範數各向異性,獲得了良好的經驗結果,但沒有解釋為什麼存在底層問題。Aurora 的分析:Muon 的極因子更新對方形矩陣做正確的事情,對高矩陣做錯誤的事情 —— 而「高」描述了大多數具有大擴展因子的 MLP,所以這個 bug 在每個人正在訓練的架構中加劇。Aurora 將權重更新重新表述為聯合約束:左半正交性和均勻行範數,同時解決而不是事後修補。兩種實現交付:黎曼 Aurora(在受約束流形上的梯度投影)和 vanilla Aurora(更簡單的實際變體)。Tilde 報告在 1.1B 規模上的開源網際網路資料上 100× 資料效率,在 modded-nanoGPT speedrun 上的新 state-of-the-art(超過 NorMuon 之前的 SOTA),並在 HellaSwag 上超過更大的模型。100× 聲明需要 harness 公開,然後才能被視為福音 —— 那是一個世代結果,不是漸進的 —— 但 speedrun SOTA 是更可驗證的點,因為它有每個人都可以比較的公共參考數字。
Muon 自 2024 年末以來一直在獲得採用,作為比 AdamW 更高效的計算替代方案,特別是對於 nanoGPT 風格的 speedruns 和越來越多地用於生產前沿訓練運行。Aurora 的診斷意味著每個從 Muon 獲得良好結果的人都在第 500 步時悄悄地損失了大約四分之一的 MLP 容量給死神經元 —— 推測訓練後期會更多。NorMuon 已經是人們感覺到有什麼不對但沒有解釋的標誌。更廣泛的模式:優化器研究有一個安靜的十年,AdamW 被視為已解決,而最近的浪潮(Lion、Sophia、Muon、NorMuon,現在是 Aurora)正在重新打開這個問題。drop-in 替代框架和 6% 計算開銷是使 Aurora 真正可採用而不是研究好奇心的部分 —— 如果它乾淨地移植到現有訓練管道,從 Muon 切換的門檻很低。死神經元數字也是任何人添加到訓練運行儀表板的有用診斷,無論他們最終選擇哪個優化器。
GitHub 程式碼在 `tilde-research/aurora-release`,論文在 blog.tilderesearch.com。如果你在 100M 參數規模以上訓練 transformers 並使用 Muon,在相信 100× 頭條數字之前,Aurora 值得在你的具體工作負載上進行受控 A/B 測試。神經元死亡框架是應該讓任何在 Muon 上運行生產訓練的人擔心的部分 —— 你可能一直在損失你不知道自己在損失的容量。對其他所有人來說,優化器研究浪潮繼續表明「訓練穩定性」和「訓練效率」背後仍有大量未解決的問題,以及解決它們的實驗室對堆疊其餘部分獲得不成比例的槓桿。