月球計畫人工智慧與加州大學洛杉磯分校研究人員發布月光：一個使用5.7T標記訓練的3B/16B參數專家混合模型，採用穆昂優化器

訓練大型語言模型（LLMs）已成為推進人工智慧的重要部分，但這並不容易。隨著模型的大小和數據集的增長，傳統的優化方法，尤其是AdamW，開始顯示出其限制。主要的挑戰之一是管理計算成本，並確保在長時間的訓練過程中保持穩定。像消失或爆炸的梯度、不同參數矩陣之間不一致的更新幅度，以及分散環境對資源的高需求，都使這個過程變得複雜。簡而言之，隨著研究人員朝著擁有數十億個參數和數萬億個標記的模型推進，迫切需要更精細的優化技術來更有效和穩定地處理這些複雜性。

為了應對這些挑戰，Moonshot AI與加州大學洛杉磯分校（UCLA）合作開發了Moonlight——一種使用Muon優化器的專家混合模型（MoE）。Moonlight有兩種配置：一種是啟用3億個參數的版本，總共擁有160億個參數，訓練於5.7萬億個標記上。這項工作基於最初為小型模型設計的Muon優化器，通過擴展其原則來滿足大型訓練的需求。Muon的核心創新在於通過牛頓-舒爾茨迭代進行矩陣正交化。這種方法有助於確保梯度更新在模型的參數空間中更均勻地應用。通過解決與AdamW相關的常見問題，Muon提供了一個有前景的替代方案，增強了訓練的效率和穩定性。

技術細節

深入了解Moonlight背後的技術創新，可以看到對Muon優化器所做的周到調整。兩個主要的修改是使Muon適合大規模訓練的關鍵。首先，整合權重衰減——這是一種通常與AdamW一起使用的技術——有助於控制權重大小的增長，特別是在使用大型模型和大量標記進行訓練時。如果沒有權重衰減，權重和層輸出可能會過度增長，從而可能隨著時間推移而降低模型性能。

第二個調整涉及校準每個參數的更新幅度。在實踐中，Muon中的更新幅度可能根據權重矩陣的形狀而有所不同。為了使這些更新一致，該方法按每個矩陣最大維度的平方根比例進行縮放。這一改變使Muon的行為與AdamW的性能更為一致，並確保所有參數都得到一致的更新。

此外，Muon的分散實現基於ZeRO-1的技術，將優化器狀態分割到數據並行組中。這種方法減少了內存開銷，並限制了通常與分散訓練相關的通信成本。雖然需要額外的步驟，例如收集梯度和執行牛頓-舒爾茨迭代，但這些步驟已經過優化，以使其對整體訓練時間的影響保持在最低限度。結果是一個保持競爭性能的優化器，同時需要更少的計算資源。

實證結果和數據分析的見解

對Moonlight的實證評估強調了這些技術改進的實際好處。在1.2萬億個標記的中間檢查點，Moonlight在語言理解的任務中顯示出比使用AdamW訓練的對應模型（稱為Moonlight-A）和其他類似的MoE模型略微更好的表現。例如，在評估語言理解的基準測試中，Moonlight在MMLU等基準上獲得了稍高的分數。在代碼生成任務中，其性能提升更加明顯，這表明Muon的精細更新機制有助於整體任務表現的改善。

擴展法則實驗進一步說明了Muon的優勢。這些實驗顯示，Muon可以在僅使用約一半的訓練計算成本的情況下，匹配使用AdamW訓練的模型的性能。這種效率對於在資源限制與推進模型能力之間取得平衡的研究人員來說，具有重要意義。此外，對權重矩陣的光譜分析表明，Moonlight使用Muon進行訓練會導致更廣泛的奇異值範圍。這種更新方向的多樣性可能有助於模型在各種任務中的泛化能力。

在監督微調階段的額外研究表明，當預訓練和微調都使用Muon進行時，這種優化器的好處在整個訓練過程中持續存在。在預訓練和微調之間切換優化器的情況下，差異不太明顯，這表明優化方法的一致性是有益的。

結論

總而言之，Moonlight的開發代表了大型語言模型訓練的一次深思熟慮的進步。通過採用Muon優化器，Moonshot AI和UCLA的團隊提供了一種可行的替代方案，顯示出在訓練效率和模型穩定性方面的改進。主要增強包括整合權重衰減和對每個參數更新幅度的調整，這兩者都有助於在不同類型的權重矩陣之間協調更新。分散實現進一步強調了這種方法的實際好處，特別是在減少大型訓練環境中的內存和通信開銷方面。

從Moonlight項目中獲得的見解在技術報告“Muon is Scalable for LLM Training”中得到了清晰的闡述。這項工作顯示，在計算最優條件下，Muon可以實現與AdamW相當甚至更優的性能，同時顯著降低計算成本。報告還強調，從AdamW轉向Muon不需要大量的超參數調整，簡化了研究人員的整合過程。

展望未來，Muon實現的開源以及預訓練模型和中間檢查點的發布，預計將促進對可擴展優化技術的進一步研究。未來的工作可能會探索將Muon擴展到其他範數約束，或將其優勢整合到涵蓋所有模型參數的統一優化框架中。這些努力可能會導致更強大和高效的訓練策略，逐步塑造大型語言模型開發的新標準。