記憶網(wǎng)絡(luò)（MemNN）是神經(jīng)網(wǎng)絡(luò)的重要分支，在問答領(lǐng)域有廣泛應(yīng)用。其計(jì)算特點(diǎn)更符合人類記憶和思考的過程，相比傳統(tǒng)RNN、LSTM等模型具有更好的長期記憶能力，最近的關(guān)注度也很高。本文首先闡述了MemNN的網(wǎng)絡(luò)特點(diǎn)，進(jìn)而分析其架構(gòu)設(shè)計(jì)中的難點(diǎn)，主要針對其存儲過大、帶寬高、稀疏計(jì)算三個方面，結(jié)合ISCA相關(guān)論文的論述，提出一些可能的解決方案和思考方向，希望對今后的AI加速器設(shè)計(jì)有所啟發(fā)。

　　memory network（MemNN），也叫memory-argumented neural network，2014年由Facebook的工程師提出。這種特殊的網(wǎng)絡(luò)有很強(qiáng)的上下文信息感知和處理能力，非常適合信息提取，問答任務(wù)等人工智能輔助領(lǐng)域。不同于傳統(tǒng)的前饋網(wǎng)絡(luò)CNN，RNN等將訓(xùn)練集壓縮成hidden state進(jìn)行存儲的方式，這類方法產(chǎn)生的記憶太小了，在壓縮過程中損失了很多有用信息。而MemNN是將所有的信息存在一個外部memory中，和inference一起聯(lián)合訓(xùn)練，得到一個能夠存儲和更新的長期記憶模塊。這樣可以最大限度的保存有用信息，下圖是一個簡單的MemNN的流程示意圖。

　　通過描述性的語言解釋下MemNN的運(yùn)行過程。首先輸入的內(nèi)容（input story sentences）經(jīng)過加工和提取儲存在外部memory中。問題提出后，會和上述內(nèi)容一起通過embedding過程產(chǎn)生question，input和output三部分，相當(dāng)于在理解問題和內(nèi)容的相關(guān)程度。然后通過inference過程，尋找和問題最相關(guān)的內(nèi)容語句，再提取語句中和問題最相近的單詞，最終產(chǎn)生答案A。在這個過程中，輸入內(nèi)容幾乎是完整的保存和使用的，而不是像RNN那樣壓縮成數(shù)據(jù)量較少的中間狀態(tài)，因此信息的完整性是比較好的。同時inference使用的是full connection連接和softmax歸一化，很容易在傳統(tǒng)的深度學(xué)習(xí)優(yōu)化算法和硬件上獲得較好的加速效果。

　　具體到細(xì)節(jié)，MemNN主要包括2種操作：embedding和inference。前者是將輸入轉(zhuǎn)化成中間狀態(tài)的計(jì)算過程，而inference是通過多層神經(jīng)網(wǎng)絡(luò)來推斷語句和問題的相關(guān)性。story經(jīng)embedding過程產(chǎn)生2個矩陣：input 和output matrix，把輸入語句轉(zhuǎn)化為內(nèi)部存儲的向量，這一步會使用通常NLP的詞向量轉(zhuǎn)化方法。問題輸入給MemNN網(wǎng)絡(luò)后，會通過inference計(jì)算得到和上述內(nèi)部向量之間的相關(guān)性，具體是三個步驟。

　　首先“input memory representation”過程會計(jì)算問題向量和input matrix的點(diǎn)積后歸一化，得到和input matrix維度一致的概率向量p，即問題和各記憶向量的相關(guān)程度。這一步的運(yùn)算是矩陣乘累加和softmax。

　　第二，“output memory representation”過程將output matrix按概率向量p進(jìn)行加權(quán)求和，得到輸出向量o，相當(dāng)于選取了相關(guān)性最高的記憶向量組合。這一步主要是矩陣點(diǎn)積。

　　最后一步“output calculation”是將輸出向量轉(zhuǎn)化為所需答案的格式，得到各單詞相對答案的概率，運(yùn)算是全連接型的矩陣乘累加。

　　上述只做行為級的描述，相關(guān)步驟的公式可以參考論文[1]。

　　經(jīng)過上述分析，我們可以總結(jié)下MemNN的計(jì)算特點(diǎn)并從架構(gòu)角度分析其可能存在的問題。首先，也是最明顯的，MemNN對輸入內(nèi)容的保存沒有經(jīng)過大幅度的壓縮，信息完整性很高，這樣在問答推理上相比RNN等壓縮模型很有優(yōu)勢，不過帶來的問題就是存儲空間會隨著內(nèi)容的增大而線性增加，對片上存儲的壓力較大。而片上存儲不足的直接后果就是內(nèi)存帶寬需求的增加。第二，運(yùn)算幾乎都是矩陣乘累加或者點(diǎn)乘，這一點(diǎn)和RNN是比較類似的，這樣在同等運(yùn)算強(qiáng)度下，對數(shù)據(jù)量的需求要高于卷積，也就是對帶寬的需求更大；第三，由于MemNN計(jì)算的特點(diǎn)是從story生成的多個向量中選擇相關(guān)性最大的產(chǎn)生答案，因此中間結(jié)果矩陣會是一個很稀疏的矩陣，只有相關(guān)性較強(qiáng)的部分才有值，其他不相關(guān)的幾乎都是0，這樣的話通常的密集運(yùn)算加速器（如TPU等）效果就不好了，需要軟件和硬件著重考慮如何進(jìn)行稀疏性的優(yōu)化。

　　可以看出，對于MemNN來說，現(xiàn)有的ASIC加速器并不能很好的優(yōu)化上述問題，因此在計(jì)算效率上應(yīng)該是沒有GPU強(qiáng)的，這也是GPU這種SIMT結(jié)構(gòu)適應(yīng)性更強(qiáng)的優(yōu)勢。如何在硬件上同時高效支持CNN和MemNN這兩種差別較大的模型，還有比較長的路要走。今年的ISCA論文中，有一篇是針對MemNN進(jìn)行優(yōu)化，盡管主要在算法層面，但可以在一定程度上給未來的通用AI加速器作為參考。

　　針對MemNN存儲較大的優(yōu)化方法利用了input memory representaion計(jì)算的可交換性。通常的計(jì)算是先乘累加在softmax，由于softmax是一個除法運(yùn)算，分母是固定的，因此可以利用乘法分配律，將矩陣乘累加拆成若干個部分單獨(dú)執(zhí)行，這樣歸一運(yùn)算和求和計(jì)算就很好的分開了，具體公式可以參考論文[2]。針對后者，我們不需要一次性的算完e^n后再和Mout做乘法，而是將其分組（N），每組內(nèi)部先完成全部的求和，這樣一次運(yùn)算對存儲的需求就只有之前的1/N了，如果正好全部放在片上cache中，那么組內(nèi)運(yùn)算就不需要進(jìn)行內(nèi)存的替換。再配合ping-pong預(yù)取下一組的數(shù)據(jù)，可以完美的將數(shù)據(jù)預(yù)取和計(jì)算并行起來，大大降低了傳統(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)化。兩者相輔相成。

　　其他兩種方法沒什么特點(diǎn)，一個是針對稀疏性提出的根據(jù)閾值丟棄運(yùn)算，這個主要在軟件層面，CPU和GPU會比較容易實(shí)現(xiàn)；另一個是針對embedding matrix是常量的特點(diǎn)，設(shè)計(jì)了一個dedicated cache專門存放input和output matrix，類似于獨(dú)立的weight cache，進(jìn)行數(shù)據(jù)分隔式存儲，減小對計(jì)算中間結(jié)果的干擾。最后使用FPGA做了一個硬件實(shí)現(xiàn)，比CPU算法高了6倍。這個加速比并不算高，主要原因是對于cache替換的優(yōu)化在CPU和FPGA上的提升是差不多的；而稀疏矩陣的加速效果FPGA甚至?xí)陀贑PU；至于embedding cache，我個人感覺作用不大。因此FPGA主要就是運(yùn)算單元數(shù)量和某些ASIC算法（softmax）上的優(yōu)勢了。因此該方案并沒有很好的利用MemNN的特征，不是一個很好的解決方案，論文中也只是一帶而過的介紹。

　　總結(jié)一下，MemNN的運(yùn)算特點(diǎn)決定了它并不能在當(dāng)前的AI硬件加速器中獲得很好的提升。這突出了TPU類ASIC加速器算法適應(yīng)狹窄的缺點(diǎn)，也幾乎是目前所有面世的AI加速器的局限。之前在“AI芯片的趨勢”一文中提到了靈活性，對于AI算法而言，更需要軟硬一體層面的緊密配合，甚至要超過在CPU和GPU這類通用處理器上的所能做的極致。在這一點(diǎn)上，目前的AI加速器設(shè)計(jì)還有很多可以改進(jìn)之處。