英偉達開源新大模型:訓練數據減少40倍,算力節省1.8倍
全球AI領導者英偉達(Nvidia)開源了最新大模型Nemotron-4-Minitron-4B和Nemotron-4-Minitron-8B。
據悉這兩個模型是基于Meta開源的Llama-3.1 8B,但英偉達使用了兩種高效的訓練方法結構化剪枝和知識蒸餾。
相比從頭訓練,每個額外模型所需的訓練token數據更少,僅需大約1000億token,最多減少40倍,算力成本可節省1.8倍。性能卻依然媲美Llama-3.1 8B、Mistral 7B、Gemma 7B等知名模型,而這些模型是在高達15萬億token數據訓練而成。
4B開源地址:https://huggingface.co/nvidia/Nemotron-4-Minitron-4B-Base
8B開源地址:https://huggingface.co/nvidia/Nemotron-4-Minitron-8B-Base
結構化剪枝方法
剪枝作為一種大模型壓縮技術,其核心思想是去除神經網絡中不重要的權重。在大模型中,每個神經元通過一系列權重與其他神經元相連,這些權重決定了信號在網絡中的傳遞強度。通過移除那些對模型輸出影響較小的權重,從而減少模型的復雜性和計算需求。
傳統的剪枝方法是隨機或根據特定策略移除權重矩陣中的單個元素。雖然這種方法能有效減小模型的體積,但對硬件的利用效率不高,因為它破壞了權重矩陣的結構,難以充分利用現代硬件的并行計算優勢。
所以,英偉達使用了一種結構化剪枝方法,保留了權重矩陣的結構,通過移除整個神經元、注意力頭或卷積濾波器等,使得剪枝后的模型仍然適合在GPU、TPU等硬件上高效運行。這不僅降低模型的內存占用和計算需求,還可能提升模型的訓練速度和推理時間,使得大模型可以在有限的資源環境中進行部署。
在實際應用中,剪枝策略的選擇依賴于模型的具體結構和優化目標。例如,剪枝神經元意味著減小神經網絡層的寬度,剪枝注意力頭則影響模型對輸入序列的理解方式,而剪枝模型的深度則直接減少模型的層數。
英偉達的研究人員在剪枝的過程中發現,在不同的剪枝軸上,比如寬度和深度,采用不同的策略會產生不同的效果。
初期單獨剪枝神經元和注意力頭比同時剪枝神經元、注意力頭和嵌入通道更有效,但經過幾次迭代的重新訓練后,這種順序可能會發生變化。而寬度剪枝在重新訓練后通常比深度剪枝有更好的效果。
知識蒸餾
使用剪枝后的模型通常需要一個重新訓練的階段以恢復和優化性能,知識蒸餾便是最佳方法之一。
其原理非常簡單,就是讓剪枝后的“學生模型”來模仿未剪枝的“教師模型”的行為,可以在使用極少量原始訓練數據的情況下,顯著提升剪枝模型的表現。
尤其是,當大模型的深度被大幅減少時,使用logit、中間狀態和嵌入的蒸餾策略更為有效;而當深度變化不大時,僅使用logit蒸餾就足夠了。
英偉達在訓練Nemotron-4-Minitron-4B和Nemotron-4-Minitron-8B時,使用的便是基于logit的知識蒸餾方法,讓學生模型的輸出概率分布模仿教師模型,從而學習到教師模型對數據的深層理解。
作為softmax層之前的輸出,logits代表了模型對每個詞匯的原始預測分數。通過最小化學生模型和教師模型logits之間的差異,學生模型能夠學習到教師模型的軟概率分布,而不僅僅是最終的預測結果,這增加了學生模型泛化能力的信息量。
除了logits層面的蒸餾,英偉達還探索了中間層特征的蒸餾,通過將教師模型的隱藏層狀態映射到學生模型,并最小化這些狀態之間的差異,進一步提高了學生模型的性能。這種方法使學生模型能夠學習到教師模型在不同層次上的表示能力,更好地捕捉數據的復雜特征。
在損失函數的選擇上,嘗試了多種損失函數,包括Kullback-Leibler散度、均方誤差和余弦相似度等。他們發現,對于深度減少的學生模型,使用logits損失和中間層損失的組合可以取得更好的效果。
根據測試數據顯示,經過結構化剪枝和知識蒸餾的Minitron-4B和Minitron-8B模型,與從頭訓練模型相比訓練數據減少了40倍,算力需求節省1.8倍,在MMLU上的評分提升16%,性能可媲美Mistral 7B、Gemma 7B和Llama-3 8B。
本文轉自AIGC開放社區 ,作者:AIGC開放社區
