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