大模型推理速度提升50%以上，還能保證少樣本學習性能！

小米大模型團隊提出SUBLLM（Subsampling-Upsampling-Bypass Large Language Model），國際AI語音大牛、開源語音識別工具Kaldi之父Daniel Povey也參與指導。

與Llama等模型相比，SUBLLM在訓練和推理速度以及降低內(nèi)存方面都有了顯著提升。

在大模型訓練中，SUBLLM的速度提高了26%，每個GPU的內(nèi)存減少了10GB。在推理中，它的速度提高了37%，每個GPU的內(nèi)存減少了1GB。

訓練和推理速度分別最高可以提高至34%和52%。

SUBLLM通過智能地選擇和處理數(shù)據(jù)，使得模型在訓練和推理時更加高效：子采樣模塊剔除不必要的信息，上采樣模塊恢復數(shù)據(jù)的完整性，而繞過模塊則加快了學習過程。

在一萬字中挑選最關鍵的五百字

目前，云端的大模型處理超長文本任務，通常需要動用多達8個GPU，這個過程不僅耗時，而且成本昂貴。如果將大模型類比于人腦，那么當前大模型的運行功率相比于人腦運行功率的100倍以上。

此前，Daniel Povey在語音識別領域提出了Zipformer，Zipformer可以用最低壓縮16倍的幀率，達到與更大模型一致甚至更高的語音識別率，完成了語音識別領域的“四兩撥千斤”。

小米集團大模型團隊嘗試將這一思路擴展至大型語言模型中，在性能不受損害的前提下，實現(xiàn)了更高效率的大模型運算。

總的來說，SUBLLM的工作原理通過引入子采樣、上采樣和旁路模塊等方式，對計算資源動態(tài)分配，從而減少了冗余的token計算負擔，加速了模型的訓練和推理過程。

能做到就像在一萬字中挑選最關鍵的五百字一樣，保留文本中必需的部分，刪減其中的冗余，從而讓大模型所需處理的文本更短。

就實現(xiàn)路徑而言，會將子采樣模塊根據(jù)token的重要性分數(shù)對其進行篩選，保留重要的token并丟棄不重要的部分。

隨后，上采樣模塊將子采樣后的序列恢復到原始長度，確保語言模型在生成token時的順序一致性。

同時，旁路模塊通過結(jié)合子采樣前后的序列，進一步提高了模型的收斂速度。這種設計不僅顯著減少了計算成本，還保持了輸入序列的語義完整性。

如果將SUBLLM理解為一個聰明的編輯，就像我們的大腦會識別要點一樣，它可以在閱讀一大段文字時快速識別出哪些詞是關鍵的，哪些詞不那么重要。SUBLLM會保留那些重要的詞匯，而忽略那些不太重要的部分，這就大大減少了需要處理的信息量。

隨后，就像我們能通過只言片語補充完整故事的來龍去脈，SUBLLM也能將精簡后的信息恢復到原有的完整度，確保整個文本在表達時的連貫與完整。在處理信息時，SUBLLM還能更加迅速地找到最佳的表達方式。

接下來具體看SUBLLM的模型結(jié)構(gòu)。

SUBLLM具體長啥樣？

前不久，谷歌Deepmind提出了mixture of depths（MoD）模型結(jié)構(gòu)，MoD使用靜態(tài)計算預算，使用每個塊的路由器選擇token進行計算，并通過對自注意力和MLP塊或殘差連接的選擇來優(yōu)化FLOP使用。

更早以前，經(jīng)典論文CoLT5使用條件路由來決定給定token是通過輕量分支還是重量分支在前饋和注意力層中傳遞，以便將更多資源分配給重要token。

與這些模型結(jié)構(gòu)類似，SUBLLM采用的原理接近于人腦對于信息的處理機制。

人腦有兩種思維模式，一種低功耗的快模式，一種高功耗的慢模式，分工明確，且兩種模式恰恰用的是同一個腦部區(qū)域。

因此，SUBLLM作者也從這一信息處理模式的角度思考了如何將大模型的算力進行合理地分配：重要的token用全部算力，相對不重要的token使用更少算力。

具體來說，SUBLLM的模型結(jié)構(gòu)是基于decoder-only的大語言模型架構(gòu)，在不改變原有模型結(jié)構(gòu)的基礎上，在一些特殊的層上進行了結(jié)構(gòu)升級。

為了管理要處理的token數(shù)量，子采樣和上采樣模塊被集成到Transformer塊之間。

首先，模型使用幾個Transformer塊處理完整序列，捕獲全面的token序列表示。

引入子采樣模塊后，這些模塊暫時去除不關鍵的token，從而減少處理所需的序列長度。

然后對縮減后的序列進行更多次的子采樣過程，也就是序列的縮減是嵌套的。序列壓縮的最高級別發(fā)生在網(wǎng)絡的最中間的Transformer塊中。

隨后，使用上采樣模塊逐步恢復序列長度。這些模塊將較短的處理序列與子采樣前的原始序列合并，將它們恢復到完整長度。

這種機制允許僅解碼器模型作為語言模型操作，按順序生成token，保證輸入和輸出序列長度相同。

此外，上采樣過程后集成了繞過連接模塊，以利用每個子采樣前的嵌入，幫助改進從子采樣到上采樣的學習過程。

隨后的實驗證實，這種方法顯著提高了收斂效率。

與LLaMA模型相比，SUBLLM在訓練和推理方面分別實現(xiàn)了26%和37%的速度提升，同時顯著降低了內(nèi)存成本，同時保持了性能。

預訓練階段、推理階段計算效率的詳細分析：

論文鏈接：https://arxiv.org/abs/2406.06571