記憶網(wǎng)絡(luò)（MemNN）是神經(jīng)網(wǎng)絡(luò)的重要分支,，在問(wèn)答領(lǐng)域有廣泛應(yīng)用。其計(jì)算特點(diǎn)更符合人類記憶和思考的過(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ǔ)的方式,，這類方法產(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是比較類似的,，這樣在同等運(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專門(mén)存放input和output matrix,，類似于獨(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類ASIC加速器算法適應(yīng)狹窄的缺點(diǎn),，也幾乎是目前所有面世的AI加速器的局限,。之前在“AI芯片的趨勢(shì)”一文中提到了靈活性，對(duì)于AI算法而言,，更需要軟硬一體層面的緊密配合,，甚至要超過(guò)在CPU和GPU這類通用處理器上的所能做的極致。在這一點(diǎn)上,，目前的AI加速器設(shè)計(jì)還有很多可以改進(jìn)之處,。