在人工智能的激烈競爭中,一場耗資百萬美元的大規模實驗正悄然改變着大語言模型的訓練方式。階躍星辰研究團隊日前發佈重磅研究成果,他們通過耗費近100萬NVIDIA H800GPU小時的算力,從零開始訓練了3,700個不同規模的模型,累計訓練了驚人的100萬億個token,揭示出一條被稱爲"Step Law"的普適性縮放規律,爲大語言模型的高效訓練提供了全新指南。
這項研究不僅僅是對超參數優化的探索,更是第一個全面考察模型最優超參在不同形狀、稀疏度和數據分佈下穩定性的工作。研究結果表明,無論模型採用何種架構設計,無論訓練數據來自何種語言或領域,Step Law都表現出令人驚歎的魯棒性,這大大增強了該工具在實際應用中的價值。
團隊訓練的3,700個模型涵蓋了不同規模、不同超參數組合、不同形狀、不同數據配比以及不同稀疏度的配置,包括MoE和Dense兩種架構。通過這些海量實驗,他們發現最優學習率隨模型參數規模與數據規模呈現冪律變化,而最優批量大小主要與數據規模相關。這一發現顛覆了業界對超參數設置的傳統認知。
圖源備註:圖片由AI生成,圖片授權服務商Midjourney
實驗數據顯示,在固定模型大小與數據規模的條件下,超參數優化的Landscape呈現明顯的凸性特徵,這意味着存在一個穩定且容易尋找的最優超參數區域。爲了驗證這一點,研究團隊構建了三維可視化空間,直觀展示了學習率與批量大小對訓練損失的影響。結果清晰地展現出"山谷"形態,凸性底端是一個相對平坦的區域,這爲實踐中的超參數調優提供了寶貴的理論依據。
爲了讓這一發現惠及整個AI社區,團隊開發並推出了一款通用的最優超參數估算工具。該工具的預測結果與通過窮舉搜索得到的全局最優超參數相比,性能差距僅爲0.09%。這意味着研究人員和工程師們可以不再依賴昂貴的網格搜索,而是直接通過這一工具獲得接近最優的超參數配置。
更讓人印象深刻的是Step Law的普適性。研究團隊從三個不同角度驗證了其適用範圍:首先,無論模型形狀如何變化——是偏向寬度、偏向深度,還是寬深平衡——Step Law都能準確預測最優超參數區域;其次,這一規律不僅適用於Dense模型,還能很好地擴展到不同稀疏度的MoE模型;最後,無論訓練數據是英語主導、中英雙語、代碼與英語混合,還是以代碼爲主的分佈,Step Law都表現出了驚人的穩定性。
研究還揭示了學習率調度策略的優化方向。與傳統的學習率衰減策略不同,團隊提出採用固定的最小學習率(1e-5),而非傳統方法中將最小值設爲最大值的十分之一。這一改變使得訓練在後期能夠維持更爲合理的參數更新步長,有效避免了損失函數在收斂階段的持續振盪。
此外,研究發現平滑訓練損失與驗證損失的最優超參數高度一致,這一發現爲超參數選擇提供了更爲經濟的方法——研究人員可以通過監控平滑訓練損失來指導超參數調整,而無需頻繁在驗證集上評估模型性能。
儘管取得了顯著成果,階躍星辰研究團隊坦言這僅僅是一個開始。他們計劃陸續開源實驗的各個細節,包括近4000個模型的最終檢查點,以供整個社區進行更深入的分析和理論解釋。未來的研究方向包括探索Loss-BS-LR三維空間的凸性、改進最優超參數的擬合方法、解釋不同配置下次優區域的變化,以及深入研究不同設置下的訓練動態。
Predictable Scale系列的後續工作可能將進一步討論超大模型性能預測、Code&Math的縮放性質,以及不同Attention類型的縮放特性。可以預見,這一系列研究將爲大語言模型的高效訓練提供更全面的理論指導和實踐工具,推動AI技術向更高效、更可控的方向發展。