谷歌更新Transformer架構(gòu)，更節(jié)省計算資源！50%性能提升
谷歌終于更新了Transformer架構(gòu)。
最新發(fā)布的Mixture-of-Depths（MoD），改變了以往Transformer計算模式。
它通過動態(tài)分配大模型中的計算資源，跳過一些不必要計算，顯著提高訓(xùn)練效率和推理速度。
結(jié)果顯示，在等效計算量和訓(xùn)練時間上，MoD每次向前傳播所需的計算量更小，而且后訓(xùn)練采樣過程中步進速度提高50%。
這一方法剛剛發(fā)布，就馬上引發(fā)關(guān)注。
MoE風(fēng)頭正盛，MoD已經(jīng)來后浪拍前浪了？
還有人開始“算賬”：
聽說GPT-4 Turbo在Blackwell上提速30倍，再加上這個方法和其他各種加速，下一代生成模型可以走多遠(yuǎn)？
所以MoD如何實現(xiàn)？
迫使大模型關(guān)注真正重要信息
這項研究提出，現(xiàn)在的大模型訓(xùn)練和推理中，有很多計算是沒必要的。
比如預(yù)測下一個句子很難，但是預(yù)測句子結(jié)束的標(biāo)點符號很簡單。如果給它們分配同樣的計算資源，那么后者明顯浪費了。
在理想情況下， 模型應(yīng)該只給需要準(zhǔn)確預(yù)測的token分配更多計算資源。
所以研究人員提出了MoD。
它在輸入序列中的特定位置動態(tài)分配FLOPs（運算次數(shù)或計算資源），優(yōu)化不同層次的模型深度中的分配。
通過限制給定層的自注意力和MLP計算的token數(shù)量，迫使神經(jīng)網(wǎng)絡(luò)學(xué)會主要關(guān)注真正重要的信息。
因為token數(shù)量是事先定義好的，所以這個過程使用一個已知張量大小的靜態(tài)計算圖，可以在時間和模型深度上動態(tài)擴展計算量。
下圖右上圖中的橙色部分，表示沒有使用全部計算資源。
這種方法在節(jié)省計算資源的同時，還能提高效率。
這些模型在等效的FLOPS和訓(xùn)練時間上與基線性能相匹配，但每次前向傳播所需的FLOP更少，并且在訓(xùn)練后采樣時提速50%。
對比來看，如果為每一個token生成一個概率分布，每個token根據(jù)最高概率被送去對應(yīng)的“專家”，可能會導(dǎo)致負(fù)載不平衡。
如果反過來，這能保障負(fù)載平衡，但是可能導(dǎo)致某些token被過度處理或處理不足。
最后來看論文中使用的Expert-choice MoD，router輸出的權(quán)重被用于確定哪些token將使用transformer虧啊計算。權(quán)重較大的token將參與計算，權(quán)重較小的token將通過殘差連接繞過計算，從而解決每次向前傳播的FLOPs。
最后，研究團隊展示了MoD在不同實驗中的性能表現(xiàn)。
首先，他們使用相對較小的FLOP預(yù)算（6e18），以確定最佳超參數(shù)配置。
通過這些實驗，作者發(fā)現(xiàn)MoD方法能夠“拉低并向右推移”isoFLOP基線曲線，這意味著最優(yōu)的MoD方法在更低的損失水平上擁有更多的參數(shù)。
通過isoFLOP分析，比較6e18、2e19和1e20 FLOPs的總計算預(yù)算下的模型性能。
結(jié)果顯示，在更多FLOP預(yù)算下，F(xiàn)LOP最優(yōu)的MoD仍然比基線模型有更多的參數(shù)。
存在一些MoD變體，在步驟速度上比isoFLOP最優(yōu)基線模型更快，同時實現(xiàn)更低的損失。這表明在訓(xùn)練之外，MoD的計算節(jié)省仍然有效。
同時，研究團隊還探討了MoD和MoE結(jié)合的可能性——MoDE。
結(jié)果表明而這結(jié)合能提供更好的性能和更快的推理速度。
網(wǎng)友：聯(lián)想到了ResNet
MoD推出后馬上引發(fā)了不小關(guān)注。
有人感慨，MoE還沒有弄清楚呢，MoD都已經(jīng)來了！
這么高效的方法，讓人馬上聯(lián)想到了ResNet。
不過和ResNet不同，MoD跳過連接是完全繞過層的。
還有人表示，希望這種方法是完全動態(tài)的，而不是每個層固定百分比。
這項研究由DeepMind和麥吉爾大學(xué)共同帶來。
主要貢獻者是David Raposo和Adam Santoro。
他們二人都是DeepMind的研究科學(xué)家。此前共同帶來了神作《Relational inductive biases, deep learning, and graph networks》。
這篇論文目前被引次數(shù)超過3500次，論文核心定義了Inductive bias（歸納偏置）概念。
論文地址：???https://arxiv.org/abs/2404.02258??
本文轉(zhuǎn)自 量子位 ，作者：量子位
