大模型推理速度提升 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ù)的完整性，而繞過模塊則加快了學習過程。
在一萬字中挑選最關(guān)鍵的五百字
目前，云端的大模型處理超長文本任務(wù)，通常需要動用多達 8 個 GPU，這個過程不僅耗時，而且成本昂貴。如果將大模型類比于人腦，那么當前大模型的運行功率相比于人腦運行功率的 100 倍以上。
此前，Daniel Povey 在語音識別領(lǐng)域提出了 Zipformer，Zipformer 可以用最低壓縮 16 倍的幀率，達到與更大模型一致甚至更高的語音識別率，完成了語音識別領(lǐng)域的“四兩撥千斤”。
小米集團大模型團隊嘗試將這一思路擴展至大型語言模型中，在性能不受損害的前提下，實現(xiàn)了更高效率的大模型運算。
總的來說，SUBLLM 的工作原理通過引入子采樣、上采樣和旁路模塊等方式，對計算資源動態(tài)分配，從而減少了冗余的 token 計算負擔，加速了模型的訓練和推理過程。
能做到就像在一萬字中挑選最關(guān)鍵的五百字一樣，保留文本中必需的部分，刪減其中的冗余，從而讓大模型所需處理的文本更短。
就實現(xiàn)路徑而言，會將子采樣模塊根據(jù) token 的重要性分數(shù)對其進行篩選，保留重要的 token 并丟棄不重要的部分。
隨后，上采樣模塊將子采樣后的序列恢復到原始長度，確保語言模型在生成 token 時的順序一致性。
同時，旁路模塊通過結(jié)合子采樣前后的序列，進一步提高了模型的收斂速度。這種設(shè)計不僅顯著減少了計算成本，還保持了輸入序列的語義完整性。
如果將 SUBLLM 理解為一個聰明的編輯，就像我們的大腦會識別要點一樣，它可以在閱讀一大段文字時快速識別出哪些詞是關(guān)鍵的，哪些詞不那么重要。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)的基礎(chǔ)上，在一些特殊的層上進行了結(jié)構(gòu)升級。
為了管理要處理的 token 數(shù)量，子采樣和上采樣模塊被集成到 Transformer 塊之間。
首先，模型使用幾個 Transformer 塊處理完整序列，捕獲全面的 token 序列表示。引入子采樣模塊后，這些模塊暫時去除不關(guān)鍵的 token，從而減少處理所需的序列長度。
然后對縮減后的序列進行更多次的子采樣過程，也就是序列的縮減是嵌套的。序列壓縮的最高級別發(fā)生在網(wǎng)絡(luò)的最中間的 Transformer 塊中。
隨后，使用上采樣模塊逐步恢復序列長度。這些模塊將較短的處理序列與子采樣前的原始序列合并，將它們恢復到完整長度。
這種機制允許僅解碼器模型作為語言模型操作，按順序生成 token，保證輸入和輸出序列長度相同。
此外，上采樣過程后集成了繞過連接模塊，以利用每個子采樣前的嵌入，幫助改進從子采樣到上采樣的學習過程。
隨后的實驗證實，這種方法顯著提高了收斂效率。
與 LLaMA 模型相比，SUBLLM 在訓練和推理方面分別實現(xiàn)了 26% 和 37% 的速度提升，同時顯著降低了內(nèi)存成本，同時保持了性能。
預訓練階段、推理階段計算效率的詳細分析：
論文鏈接：https://arxiv.org/abs/2406.06571
廣告聲明：文內(nèi)含有的對外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。