谷歌終于更新了Transformer架構(gòu)。

最新發(fā)布的Mixture-of-Depths（MoD），改變了以往Transformer計(jì)算模式。

它通過(guò)動(dòng)態(tài)分配大模型中的計(jì)算資源，跳過(guò)一些不必要計(jì)算，顯著提高訓(xùn)練效率和推理速度。

結(jié)果顯示，在等效計(jì)算量和訓(xùn)練時(shí)間上，MoD每次向前傳播所需的計(jì)算量更小，而且后訓(xùn)練采樣過(guò)程中步進(jìn)速度提高50%。

這一方法剛剛發(fā)布，就馬上引發(fā)關(guān)注。

MoE風(fēng)頭正盛，MoD已經(jīng)來(lái)后浪拍前浪了？

所以MoD如何實(shí)現(xiàn)？

迫使大模型關(guān)注真正重要信息

這項(xiàng)研究提出，現(xiàn)在的大模型訓(xùn)練和推理中，有很多計(jì)算是沒(méi)必要的。

比如預(yù)測(cè)下一個(gè)句子很難，但是預(yù)測(cè)句子結(jié)束的標(biāo)點(diǎn)符號(hào)很簡(jiǎn)單。如果給它們分配同樣的計(jì)算資源，那么后者明顯浪費(fèi)了。

在理想情況下， 模型應(yīng)該只給需要準(zhǔn)確預(yù)測(cè)的token分配更多計(jì)算資源。

所以研究人員提出了MoD。

它在輸入序列中的特定位置動(dòng)態(tài)分配FLOPs（運(yùn)算次數(shù)或計(jì)算資源），優(yōu)化不同層次的模型深度中的分配。

通過(guò)限制給定層的自注意力和MLP計(jì)算的token數(shù)量，迫使神經(jīng)網(wǎng)絡(luò)學(xué)會(huì)主要關(guān)注真正重要的信息。

因?yàn)閠oken數(shù)量是事先定義好的，所以這個(gè)過(guò)程使用一個(gè)已知張量大小的靜態(tài)計(jì)算圖，可以在時(shí)間和模型深度上動(dòng)態(tài)擴(kuò)展計(jì)算量。

下圖右上圖中的橙色部分，表示沒(méi)有使用全部計(jì)算資源。

這種方法在節(jié)省計(jì)算資源的同時(shí)，還能提高效率。

這些模型在等效的FLOPS和訓(xùn)練時(shí)間上與基線性能相匹配，但每次前向傳播所需的FLOP更少，并且在訓(xùn)練后采樣時(shí)提速50%。

對(duì)比來(lái)看，如果為每一個(gè)token生成一個(gè)概率分布，每個(gè)token根據(jù)最高概率被送去對(duì)應(yīng)的“專(zhuān)家”，可能會(huì)導(dǎo)致負(fù)載不平衡。

如果反過(guò)來(lái)，這能保障負(fù)載平衡，但是可能導(dǎo)致某些token被過(guò)度處理或處理不足。

最后來(lái)看論文中使用的Expert-choice MoD，router輸出的權(quán)重被用于確定哪些token將使用transformer虧啊計(jì)算。權(quán)重較大的token將參與計(jì)算，權(quán)重較小的token將通過(guò)殘差連接繞過(guò)計(jì)算，從而解決每次向前傳播的FLOPs。

最后，研究團(tuán)隊(duì)展示了MoD在不同實(shí)驗(yàn)中的性能表現(xiàn)。

首先，他們使用相對(duì)較小的FLOP預(yù)算（6e18），以確定最佳超參數(shù)配置。

通過(guò)這些實(shí)驗(yàn)，作者發(fā)現(xiàn)MoD方法能夠“拉低并向右推移”isoFLOP基線曲線，這意味著最優(yōu)的MoD方法在更低的損失水平上擁有更多的參數(shù)。

通過(guò)isoFLOP分析，比較6e18、2e19和1e20 FLOPs的總計(jì)算預(yù)算下的模型性能。

結(jié)果顯示，在更多FLOP預(yù)算下，F(xiàn)LOP最優(yōu)的MoD仍然比基線模型有更多的參數(shù)。

存在一些MoD變體，在步驟速度上比isoFLOP最優(yōu)基線模型更快，同時(shí)實(shí)現(xiàn)更低的損失。這表明在訓(xùn)練之外，MoD的計(jì)算節(jié)省仍然有效。

同時(shí)，研究團(tuán)隊(duì)還探討了MoD和MoE結(jié)合的可能性——MoDE。

結(jié)果表明而這結(jié)合能提供更好的性能和更快的推理速度。

網(wǎng)友：聯(lián)想到了ResNet

MoD推出后馬上引發(fā)了不小關(guān)注。

有人感慨，MoE還沒(méi)有弄清楚呢，MoD都已經(jīng)來(lái)了！

這么高效的方法，讓人馬上聯(lián)想到了ResNet。

不過(guò)和ResNet不同，MoD跳過(guò)連接是完全繞過(guò)層的。

還有人表示，希望這種方法是完全動(dòng)態(tài)的，而不是每個(gè)層固定百分比。

這項(xiàng)研究由DeepMind和麥吉爾大學(xué)共同帶來(lái)。

主要貢獻(xiàn)者是David Raposo和Adam Santoro。

他們二人都是DeepMind的研究科學(xué)家。此前共同帶來(lái)了神作《Relational inductive biases, deep learning, and graph networks》。

這篇論文目前被引次數(shù)超過(guò)3500次，論文核心定義了Inductive bias（歸納偏置）概念。