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