本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯系出處。

我們常常被教育的做事“三思而后行”，充分運用積累過的經驗，現在這句話對AI也有所啟發了。

傳統的決策AI模型由于遺忘效應的存在不能有效積累經驗，但一項由華人主導的研究改變了AI的記憶方式。

新的記憶方式模仿了人類大腦，有效地提高了AI積累經驗的效率，從而將AI打游戲的成績提高了29.9%。

研究團隊由六人組成，分別來自米拉-魁北克AI研究院和微軟蒙特利爾研究所，其中有四名是華人。

他們將成果命名為的帶有記憶的決策Transformer(DT-Mem)。

相比傳統的決策模型，DT-Mem適用廣泛性更強，模型運算的效率也更高。

除了應用效果，DT-Mem的訓練時間也從最短200小時縮短至50小時。

同時，團隊還提出了一種微調方式，讓DT-Mem能夠適應未訓練過的新場景。

微調后的模型，面對沒有學習過的游戲，也能擁有不錯的表現。

工作機制受到人類啟發

傳統的決策模型基于LLM進行設計，采用隱性記憶，其性能依賴于數據和計算。

隱性記憶是無意識產生而非刻意記住的，因而也無法有意識地進行調用。

說得通俗一些，就是明明有關內容就存儲在那里，但是模型卻并不知道它的存在。

隱性記憶的這一特點決定了傳統模型存在遺忘現象，導致其工作效率往往不高。

遺忘現象表現為，在學了新的問題解決方式之后，模型可能會將舊的內容忘記，哪怕新舊問題是同一類型。

而人腦采用分布式記憶存儲方式，記憶的內容被分散存儲在大腦中的多個不同區域。

這種方式有助于有效地管理和組織多種技能，從而減輕遺忘現象。

受此啟發，研究團隊提出了一個內部工作記憶模塊來存儲、混合和檢索不同下游任務的信息。

具體而言，DT-Mem由Transformer、記憶模塊和多層感知(MLP)模塊三部分組成。

DT-Mem的Transformer模仿了GPT-2的架構，但刪去了注意力機制后的前饋層。

同時，GPT-2中的MLP模塊被拆分成了獨立組件，作為DT-Mem的一部分。

在二者之間，研究團隊引入了一個工作記憶模塊，用于存儲和處理中間信息。

這一結構是受到神經圖靈機的啟發，其中的記憶被用于推斷多種算法。

記憶模塊分析Transformer輸出的信息，并決定其存儲位置以及與已有信息的整合方式。

此外，該模塊還要考慮這些信息在今后做出決策的過程當中如何使用。

這些任務大概通過五個步驟來完成，記憶模塊首先被初始化為一個隨機矩陣。

然后是對輸入信息的整理，這一步并不是將信息傳給Transformer，而是以元組形式存入同一空間。

之后就要確定存儲位置。人類通常會將相關的信息存儲到同一位置，DT-Mem也是基于這一原理。

最后兩步——記憶更新和檢索是記憶模塊的核心，也是整個DT-Mem中最重要的環節。

記憶更新，即對已有信息進行編輯替換，以確保信息能根據任務需要及時更新。

這一步中DT-Mem會計算擦除和寫入兩個向量，進而判斷如何與已有數據混合。

記憶檢索則是對已有信息的訪問和恢復，在需要做出決策時及時調取相關有用信息。

投入實際使用之前，DT-Mem還要經歷預訓練過程。

而對于DT-Mem的微調，團隊也提出了一種新的方式。

由于使用的是基于任務進行標記的數據，這種微調能夠幫助DT-Mem適應新的任務。

這一過程基于低秩適應(LoRA)進行，在已有的矩陣中加入低秩元素。

訓練時間最多縮短32倍

為了測試DT-Mem的決策能力，研究團隊讓它玩了幾款游戲。

游戲一共有5款，全部來自Atari公司。

同時，團隊還測試了傳統模型M[ulti-game]DT的表現作為參照。

結果，DT-Mem在其中4款游戲里的最好成績均勝過MDT。

具體而言，DT-Mem比MDT的DQN標準化分數提高了29.9%。

但是，DT-Mem的參數量只有20M，僅是MDT（200M參數）的10%。

這樣的表現，說是四兩撥千斤一點也不過分。

除了表現優異，DT-Mem的訓練效率也完爆MDT。

13M參數量版本的MDT需要200小時進行訓練，而20M的DT-Mem卻只需要50個小時。

如果和200M的版本相比，訓練時間足足縮短了32倍，表現卻更優異。

而針對團隊提出的微調方式的測試結果也表明，這種微調增強了DT-Mem適應未知情景的能力。

需要說明的是，下表中用來測試的游戲對于MDT來說是已知的，因此MDT的表現在這一輪當中不作為衡量依據。

除了玩游戲，團隊還使用了Meta-World ML45基準對DT-Mem進行了測試。

這次用作參照的是H[yper]DT和P[romot]DT。

結果顯示，未經微調的模型當中，DT-Mem成績比HDT高出8個百分點。

需要說明的是，這里測試的HDT本身參數量雖然只有69K，但需依賴于2.3M參數量的預訓練模型，因此實際的參數量是DT-Mem(147K)的10余倍。

論文地址：https://arxiv.org/abs/2305.16338