傳說(shuō)中GPT-4的“致勝法寶”——MoE（混合專家）架構(gòu)，自己也能手搓了！

Hugging Face上有一位機(jī)器學(xué)習(xí)大神，分享了如何從頭開(kāi)始建立一套完整的MoE系統(tǒng)。

這個(gè)項(xiàng)目被作者叫做MakeMoE，詳細(xì)講述了從注意力構(gòu)建到形成完整MoE模型的過(guò)程。

作者介紹，MakeMoE是受到OpenAI創(chuàng)始成員Andrej Karpathy的makemore啟發(fā)并以之為基礎(chǔ)編寫(xiě)的。

makemore是一個(gè)針對(duì)自然語(yǔ)言處理和機(jī)器學(xué)習(xí)的教學(xué)項(xiàng)目，意在幫助學(xué)習(xí)者理解并實(shí)現(xiàn)一些基本模型。

同樣，MakeMoE也是在一步步的搭建過(guò)程中，幫助學(xué)習(xí)者更深刻地理解混合專家模型。

那么，這份“手搓攻略”具體都講了些什么呢？

從頭開(kāi)始搭建MoE模型

和Karpathy的makemore相比，MakeMoE用稀疏的專家混合體代替了孤立的前饋神經(jīng)網(wǎng)絡(luò)，同時(shí)加入了必要的門(mén)控邏輯。

同時(shí)，由于過(guò)程中需要用到ReLU激活函數(shù)，makemore中的默認(rèn)初始化方式被替換成了Kaiming He方法。

想要?jiǎng)?chuàng)建一個(gè)MoE模型，首先要理解自注意力機(jī)制。

模型首先通過(guò)線性變換，將輸入序列變換成用查詢（Q）、鍵（K）和值（V）表示的參數(shù)。

這些參數(shù)隨后被用于計(jì)算注意力分?jǐn)?shù)，這些分?jǐn)?shù)決定了在生成每個(gè)token時(shí)，模型對(duì)序列中每個(gè)位置的關(guān)注程度。

為了確保模型在生成文本時(shí)的自回歸特性，即只能基于已經(jīng)生成的token來(lái)預(yù)測(cè)下一個(gè)token，作者使用了多頭因果自注意力機(jī)制。

這種機(jī)制通過(guò)一個(gè)掩碼來(lái)實(shí)現(xiàn)將未處理的位置的注意力分?jǐn)?shù)設(shè)置為負(fù)無(wú)窮大，這樣這些位置的權(quán)重就會(huì)變?yōu)榱恪?/p>

多頭因果則是讓模型并行地執(zhí)行多個(gè)這樣的注意力計(jì)算，每個(gè)頭關(guān)注序列的不同部分。

完成自注意力機(jī)制的配置后，就可以創(chuàng)建專家模塊了，這里的“專家模塊”是一種多層感知器。

每個(gè)專家模塊包含一個(gè)線性層，它將嵌入向量映射到一個(gè)更大的維度，然后通過(guò)非線性激活函數(shù)（如ReLU），再通過(guò)另一個(gè)線性層將向量映射回原始的嵌入維度。

這樣的設(shè)計(jì)使得每個(gè)專家能夠?qū)Ｗ⒂谔幚磔斎胄蛄械牟煌糠郑⑼ㄟ^(guò)門(mén)控網(wǎng)絡(luò)來(lái)決定在生成每個(gè)token時(shí)應(yīng)該激活哪些專家。

于是，接下來(lái)就要開(kāi)始構(gòu)建分配和管理專家的組件——門(mén)控網(wǎng)絡(luò)。

這里的門(mén)控網(wǎng)絡(luò)同樣是通過(guò)一個(gè)線性層實(shí)現(xiàn)，該層將自注意力層的輸出映射到專家模塊的數(shù)量。

這個(gè)線性層的輸出是一個(gè)分?jǐn)?shù)向量，每個(gè)分?jǐn)?shù)代表了對(duì)應(yīng)專家模塊對(duì)于當(dāng)前處理的token的重要性。

門(mén)控網(wǎng)絡(luò)會(huì)計(jì)算這個(gè)分?jǐn)?shù)向量的top-k值并記錄其索引，然后從中選擇top-k個(gè)最大的分?jǐn)?shù)，用來(lái)加權(quán)對(duì)應(yīng)的專家模塊輸出。

為了在訓(xùn)練過(guò)程中增加模型的探索性，作者還引入了噪聲，避免所有token都傾向于被相同的專家處理。

這種噪聲通常通過(guò)在分?jǐn)?shù)向量上添加隨機(jī)的高斯噪聲實(shí)現(xiàn)。

獲得結(jié)果后，模型有選擇地將前k個(gè)值與相應(yīng)token的前k個(gè)專家的輸出相乘，然后相加形成加權(quán)和，構(gòu)成模型的輸出。

最后，將這些模塊在一起，就得到一個(gè)MoE模型了。

針對(duì)以上的整個(gè)過(guò)程，作者都提供了相應(yīng)的代碼，可以到原文中具體了解。

另外，作者還制作了端到端的Jupyter筆記，可以在學(xué)習(xí)各模塊的同時(shí)直接運(yùn)行。

感興趣的話，就趕快學(xué)起來(lái)吧！

原文地址：https://huggingface.co/blog/AviSoori1x/makemoe-from-scratch
筆記版本（GitHub）：https://github.com/AviSoori1x/makeMoE/tree/main