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