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