記憶網(wǎng)絡(luò)（MemNN）是神經(jīng)網(wǎng)絡(luò)的重要分支，在問(wèn)答領(lǐng)域有廣泛應(yīng)用。其計(jì)算特點(diǎn)更符合人類(lèi)記憶和思考的過(guò)程，相比傳統(tǒng)RNN、LSTM等模型具有更好的長(zhǎng)期記憶能力，最近的關(guān)注度也很高。本文首先闡述了MemNN的網(wǎng)絡(luò)特點(diǎn)，進(jìn)而分析其架構(gòu)設(shè)計(jì)中的難點(diǎn)，主要針對(duì)其存儲(chǔ)過(guò)大、帶寬高、稀疏計(jì)算三個(gè)方面，結(jié)合ISCA相關(guān)論文的論述，提出一些可能的解決方案和思考方向，希望對(duì)今后的AI加速器設(shè)計(jì)有所啟發(fā)。
memory network（MemNN），也叫memory-argumented neural network，2014年由Facebook的工程師提出。這種特殊的網(wǎng)絡(luò)有很強(qiáng)的上下文信息感知和處理能力，非常適合信息提取，問(wèn)答任務(wù)等人工智能輔助領(lǐng)域。不同于傳統(tǒng)的前饋網(wǎng)絡(luò)CNN，RNN等將訓(xùn)練集壓縮成hidden state進(jìn)行存儲(chǔ)的方式，這類(lèi)方法產(chǎn)生的記憶太小了，在壓縮過(guò)程中損失了很多有用信息。而MemNN是將所有的信息存在一個(gè)外部memory中，和inference一起聯(lián)合訓(xùn)練，得到一個(gè)能夠存儲(chǔ)和更新的長(zhǎng)期記憶模塊。這樣可以最大限度的保存有用信息，下圖是一個(gè)簡(jiǎn)單的MemNN的流程示意圖。
通過(guò)描述性的語(yǔ)言解釋下MemNN的運(yùn)行過(guò)程。首先輸入的內(nèi)容（input story sentences）經(jīng)過(guò)加工和提取儲(chǔ)存在外部memory中。問(wèn)題提出后，會(huì)和上述內(nèi)容一起通過(guò)embedding過(guò)程產(chǎn)生question，input和output三部分，相當(dāng)于在理解問(wèn)題和內(nèi)容的相關(guān)程度。然后通過(guò)inference過(guò)程，尋找和問(wèn)題最相關(guān)的內(nèi)容語(yǔ)句，再提取語(yǔ)句中和問(wèn)題最相近的單詞，最終產(chǎn)生答案A。在這個(gè)過(guò)程中，輸入內(nèi)容幾乎是完整的保存和使用的，而不是像RNN那樣壓縮成數(shù)據(jù)量較少的中間狀態(tài)，因此信息的完整性是比較好的。同時(shí)inference使用的是full connection連接和softmax歸一化，很容易在傳統(tǒng)的深度學(xué)習(xí)優(yōu)化算法和硬件上獲得較好的加速效果。
具體到細(xì)節(jié)，MemNN主要包括2種操作：embedding和inference。前者是將輸入轉(zhuǎn)化成中間狀態(tài)的計(jì)算過(guò)程，而inference是通過(guò)多層神經(jīng)網(wǎng)絡(luò)來(lái)推斷語(yǔ)句和問(wèn)題的相關(guān)性。story經(jīng)embedding過(guò)程產(chǎn)生2個(gè)矩陣：input 和output matrix，把輸入語(yǔ)句轉(zhuǎn)化為內(nèi)部存儲(chǔ)的向量，這一步會(huì)使用通常NLP的詞向量轉(zhuǎn)化方法。問(wèn)題輸入給MemNN網(wǎng)絡(luò)后，會(huì)通過(guò)inference計(jì)算得到和上述內(nèi)部向量之間的相關(guān)性，具體是三個(gè)步驟。
首先“input memory representation”過(guò)程會(huì)計(jì)算問(wèn)題向量和input matrix的點(diǎn)積后歸一化，得到和input matrix維度一致的概率向量p，即問(wèn)題和各記憶向量的相關(guān)程度。這一步的運(yùn)算是矩陣乘累加和softmax。
第二，“output memory representation”過(guò)程將output matrix按概率向量p進(jìn)行加權(quán)求和，得到輸出向量o，相當(dāng)于選取了相關(guān)性最高的記憶向量組合。這一步主要是矩陣點(diǎn)積。
最后一步“output calculation”是將輸出向量轉(zhuǎn)化為所需答案的格式，得到各單詞相對(duì)答案的概率，運(yùn)算是全連接型的矩陣乘累加。
上述只做行為級(jí)的描述，相關(guān)步驟的公式可以參考論文[1]。
經(jīng)過(guò)上述分析，我們可以總結(jié)下MemNN的計(jì)算特點(diǎn)并從架構(gòu)角度分析其可能存在的問(wèn)題。首先，也是最明顯的，MemNN對(duì)輸入內(nèi)容的保存沒(méi)有經(jīng)過(guò)大幅度的壓縮，信息完整性很高，這樣在問(wèn)答推理上相比RNN等壓縮模型很有優(yōu)勢(shì)，不過(guò)帶來(lái)的問(wèn)題就是存儲(chǔ)空間會(huì)隨著內(nèi)容的增大而線性增加，對(duì)片上存儲(chǔ)的壓力較大。而片上存儲(chǔ)不足的直接后果就是內(nèi)存帶寬需求的增加。第二，運(yùn)算幾乎都是矩陣乘累加或者點(diǎn)乘，這一點(diǎn)和RNN是比較類(lèi)似的，這樣在同等運(yùn)算強(qiáng)度下，對(duì)數(shù)據(jù)量的需求要高于卷積，也就是對(duì)帶寬的需求更大；第三，由于MemNN計(jì)算的特點(diǎn)是從story生成的多個(gè)向量中選擇相關(guān)性最大的產(chǎn)生答案，因此中間結(jié)果矩陣會(huì)是一個(gè)很稀疏的矩陣，只有相關(guān)性較強(qiáng)的部分才有值，其他不相關(guān)的幾乎都是0，這樣的話通常的密集運(yùn)算加速器（如TPU等）效果就不好了，需要軟件和硬件著重考慮如何進(jìn)行稀疏性的優(yōu)化。
可以看出，對(duì)于MemNN來(lái)說(shuō)，現(xiàn)有的ASIC加速器并不能很好的優(yōu)化上述問(wèn)題，因此在計(jì)算效率上應(yīng)該是沒(méi)有GPU強(qiáng)的，這也是GPU這種SIMT結(jié)構(gòu)適應(yīng)性更強(qiáng)的優(yōu)勢(shì)。如何在硬件上同時(shí)高效支持CNN和MemNN這兩種差別較大的模型，還有比較長(zhǎng)的路要走。今年的ISCA論文中，有一篇是針對(duì)MemNN進(jìn)行優(yōu)化，盡管主要在算法層面，但可以在一定程度上給未來(lái)的通用AI加速器作為參考。
針對(duì)MemNN存儲(chǔ)較大的優(yōu)化方法利用了input memory representaion計(jì)算的可交換性。通常的計(jì)算是先乘累加在softmax，由于softmax是一個(gè)除法運(yùn)算，分母是固定的，因此可以利用乘法分配律，將矩陣乘累加拆成若干個(gè)部分單獨(dú)執(zhí)行，這樣歸一運(yùn)算和求和計(jì)算就很好的分開(kāi)了，具體公式可以參考論文[2]。針對(duì)后者，我們不需要一次性的算完e^n后再和Mout做乘法，而是將其分組（N），每組內(nèi)部先完成全部的求和，這樣一次運(yùn)算對(duì)存儲(chǔ)的需求就只有之前的1/N了，如果正好全部放在片上cache中，那么組內(nèi)運(yùn)算就不需要進(jìn)行內(nèi)存的替換。再配合ping-pong預(yù)取下一組的數(shù)據(jù)，可以完美的將數(shù)據(jù)預(yù)取和計(jì)算并行起來(lái)，大大降低了傳統(tǒng)計(jì)算中數(shù)據(jù)反復(fù)替換的損失。這一點(diǎn)很有啟發(fā)，AI加速是軟硬一體的緊密結(jié)合體，算法設(shè)計(jì)中應(yīng)該充分利用數(shù)據(jù)的局部性，盡可能將一組數(shù)據(jù)和計(jì)算都在片上完成，組與組之間提高并行性。硬件應(yīng)提供相應(yīng)的并行化設(shè)計(jì)和調(diào)度方式，方便軟件進(jìn)行優(yōu)化。兩者相輔相成。
其他兩種方法沒(méi)什么特點(diǎn)，一個(gè)是針對(duì)稀疏性提出的根據(jù)閾值丟棄運(yùn)算，這個(gè)主要在軟件層面，CPU和GPU會(huì)比較容易實(shí)現(xiàn)；另一個(gè)是針對(duì)embedding matrix是常量的特點(diǎn)，設(shè)計(jì)了一個(gè)dedicated cache專(zhuān)門(mén)存放input和output matrix，類(lèi)似于獨(dú)立的weight cache，進(jìn)行數(shù)據(jù)分隔式存儲(chǔ)，減小對(duì)計(jì)算中間結(jié)果的干擾。最后使用FPGA做了一個(gè)硬件實(shí)現(xiàn)，比CPU算法高了6倍。這個(gè)加速比并不算高，主要原因是對(duì)于cache替換的優(yōu)化在CPU和FPGA上的提升是差不多的；而稀疏矩陣的加速效果FPGA甚至?xí)陀贑PU；至于embedding cache，我個(gè)人感覺(jué)作用不大。因此FPGA主要就是運(yùn)算單元數(shù)量和某些ASIC算法（softmax）上的優(yōu)勢(shì)了。因此該方案并沒(méi)有很好的利用MemNN的特征，不是一個(gè)很好的解決方案，論文中也只是一帶而過(guò)的介紹。
總結(jié)一下，MemNN的運(yùn)算特點(diǎn)決定了它并不能在當(dāng)前的AI硬件加速器中獲得很好的提升。這突出了TPU類(lèi)ASIC加速器算法適應(yīng)狹窄的缺點(diǎn)，也幾乎是目前所有面世的AI加速器的局限。之前在“AI芯片的趨勢(shì)”一文中提到了靈活性，對(duì)于AI算法而言，更需要軟硬一體層面的緊密配合，甚至要超過(guò)在CPU和GPU這類(lèi)通用處理器上的所能做的極致。在這一點(diǎn)上，目前的AI加速器設(shè)計(jì)還有很多可以改進(jìn)之處。