Meta等最新研究：多token預(yù)測，提升大模型推理效率
GPT-4、Gemini系列、Llama-3等開閉源大模型，通常使用的是下一個token預(yù)測（Next-token Prediction）的損失函數(shù)進(jìn)行預(yù)訓(xùn)練。
這種方法雖然強大，但有很多局限性，例如，需要大量的訓(xùn)練數(shù)據(jù)才能使模型達(dá)到人類兒童的智商，并且隨著模型參數(shù)的增大推理效率會變差。
因此，Meta、巴黎理工大學(xué)和巴黎薩克雷大學(xué)提出了一種全新訓(xùn)練方法多token預(yù)測（Multi-token Prediction），在訓(xùn)練的過程中要求模型在每個位置上同時預(yù)測接下來的n個Token，以提升模型推理效率，并且不會增加預(yù)訓(xùn)練時間。
研究人員在130億、67億、30億等多種不同參數(shù)的模型對該技術(shù)進(jìn)行了綜合評估。結(jié)果顯示，130億參數(shù)模型在 HumanEval上解決問題能力提高了12%，在 MBPP上解決能力提高了17%，并且推理效率也更好。
論文地址：https://arxiv.org/abs/2404.19737
多token預(yù)測架構(gòu)介紹
?
為了有效實現(xiàn)多Token預(yù)測,研究人員設(shè)計了一種巧妙的模型架構(gòu)。該架構(gòu)包含一個共享的Transformer主干網(wǎng)絡(luò),用于從輸入獲取上下文表示。
然后該上下文表示被并行輸入到n個獨立的輸出頭網(wǎng)絡(luò)中,每個輸出頭負(fù)責(zé)預(yù)測一個未來Token。在推理階段,只需使用單個下一Token預(yù)測，輸出頭即可進(jìn)行自回歸生成。而其他輸出頭則可被用于加速模型的推理效率。
對于訓(xùn)練語料中的每個位置，模型需要使用獨立的輸出頭預(yù)測接下來的n個Token。將多Token預(yù)測作為輔助訓(xùn)練任務(wù)，可以提高模型在代碼和自然語言文本方面的任務(wù)性能，而不會增加訓(xùn)練時間。
降低GPU內(nèi)存使用
?
為了解決多token預(yù)測可能導(dǎo)致GPU內(nèi)存使用量增加的問題，研究人員開發(fā)了一種前向和后向傳播順序，模型能夠減少在內(nèi)存中同時存儲的梯度數(shù)量，從而降低了內(nèi)存使用量使得訓(xùn)練更加高效。
在前向傳播過程中，模型會首先通過共享主干生成潛在表示，然后按順序計算每個獨立輸出頭的前向傳播。對于每個輸出頭，計算完畢后立即進(jìn)行后向傳播，并釋放該頭的中間數(shù)據(jù)，而不是等到所有輸出頭的前向傳播完成后才進(jìn)行。
在每個輸出頭的后向傳播中，累積梯度到共享主干，而不是在所有輸出頭計算完畢后才進(jìn)行。這樣可以確保在任何時候，內(nèi)存中只存在一個輸出頭的梯度。
優(yōu)化推理效率
?
研究人員發(fā)現(xiàn)，將多token預(yù)測與自推測解碼相結(jié)合，可以進(jìn)一步提升大模型的推理效率。與傳統(tǒng)逐個token解碼不同的是，自推測解碼允許模型一次性生成多個token，然后利用額外的輸出頭并行驗證和優(yōu)化這些預(yù)測。
這種方法顯著減少了模型生成文本所需的步驟，從而加快了模型的整體推理效率并減少了對算力的消耗。
研究人員在不同參數(shù)的模型實驗了該優(yōu)化效果，結(jié)果顯示，比傳統(tǒng)的優(yōu)化推理效率提升了3倍左右。
本文轉(zhuǎn)自 AIGC開放社區(qū) ，作者： AIGC開放社區(qū)
