本文經(jīng)AI新媒體量子位（公眾號(hào)ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

我們常常被教育的做事“三思而后行”，充分運(yùn)用積累過的經(jīng)驗(yàn)，現(xiàn)在這句話對(duì)AI也有所啟發(fā)了。

傳統(tǒng)的決策AI模型由于遺忘效應(yīng)的存在不能有效積累經(jīng)驗(yàn)，但一項(xiàng)由華人主導(dǎo)的研究改變了AI的記憶方式。

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

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

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

相比傳統(tǒng)的決策模型，DT-Mem適用廣泛性更強(qiáng)，模型運(yùn)算的效率也更高。

除了應(yīng)用效果，DT-Mem的訓(xùn)練時(shí)間也從最短200小時(shí)縮短至50小時(shí)。

同時(shí)，團(tuán)隊(duì)還提出了一種微調(diào)方式，讓DT-Mem能夠適應(yīng)未訓(xùn)練過的新場(chǎng)景。

微調(diào)后的模型，面對(duì)沒有學(xué)習(xí)過的游戲，也能擁有不錯(cuò)的表現(xiàn)。

工作機(jī)制受到人類啟發(fā)

傳統(tǒng)的決策模型基于LLM進(jìn)行設(shè)計(jì)，采用隱性記憶，其性能依賴于數(shù)據(jù)和計(jì)算。

隱性記憶是無意識(shí)產(chǎn)生而非刻意記住的，因而也無法有意識(shí)地進(jìn)行調(diào)用。

說得通俗一些，就是明明有關(guān)內(nèi)容就存儲(chǔ)在那里，但是模型卻并不知道它的存在。

隱性記憶的這一特點(diǎn)決定了傳統(tǒng)模型存在遺忘現(xiàn)象，導(dǎo)致其工作效率往往不高。

遺忘現(xiàn)象表現(xiàn)為，在學(xué)了新的問題解決方式之后，模型可能會(huì)將舊的內(nèi)容忘記，哪怕新舊問題是同一類型。

而人腦采用分布式記憶存儲(chǔ)方式，記憶的內(nèi)容被分散存儲(chǔ)在大腦中的多個(gè)不同區(qū)域。

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

受此啟發(fā)，研究團(tuán)隊(duì)提出了一個(gè)內(nèi)部工作記憶模塊來存儲(chǔ)、混合和檢索不同下游任務(wù)的信息。

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

DT-Mem的Transformer模仿了GPT-2的架構(gòu)，但刪去了注意力機(jī)制后的前饋層。

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

在二者之間，研究團(tuán)隊(duì)引入了一個(gè)工作記憶模塊，用于存儲(chǔ)和處理中間信息。

這一結(jié)構(gòu)是受到神經(jīng)圖靈機(jī)的啟發(fā)，其中的記憶被用于推斷多種算法。

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

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

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

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

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

最后兩步——記憶更新和檢索是記憶模塊的核心，也是整個(gè)DT-Mem中最重要的環(huán)節(jié)。

記憶更新，即對(duì)已有信息進(jìn)行編輯替換，以確保信息能根據(jù)任務(wù)需要及時(shí)更新。

這一步中DT-Mem會(huì)計(jì)算擦除和寫入兩個(gè)向量，進(jìn)而判斷如何與已有數(shù)據(jù)混合。

記憶檢索則是對(duì)已有信息的訪問和恢復(fù)，在需要做出決策時(shí)及時(shí)調(diào)取相關(guān)有用信息。

投入實(shí)際使用之前，DT-Mem還要經(jīng)歷預(yù)訓(xùn)練過程。

而對(duì)于DT-Mem的微調(diào)，團(tuán)隊(duì)也提出了一種新的方式。

由于使用的是基于任務(wù)進(jìn)行標(biāo)記的數(shù)據(jù)，這種微調(diào)能夠幫助DT-Mem適應(yīng)新的任務(wù)。

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

訓(xùn)練時(shí)間最多縮短32倍

為了測(cè)試DT-Mem的決策能力，研究團(tuán)隊(duì)讓它玩了幾款游戲。

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

同時(shí)，團(tuán)隊(duì)還測(cè)試了傳統(tǒng)模型M[ulti-game]DT的表現(xiàn)作為參照。

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

具體而言，DT-Mem比MDT的DQN標(biāo)準(zhǔn)化分?jǐn)?shù)提高了29.9%。

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

這樣的表現(xiàn)，說是四兩撥千斤一點(diǎn)也不過分。

除了表現(xiàn)優(yōu)異，DT-Mem的訓(xùn)練效率也完爆MDT。

13M參數(shù)量版本的MDT需要200小時(shí)進(jìn)行訓(xùn)練，而20M的DT-Mem卻只需要50個(gè)小時(shí)。

如果和200M的版本相比，訓(xùn)練時(shí)間足足縮短了32倍，表現(xiàn)卻更優(yōu)異。

而針對(duì)團(tuán)隊(duì)提出的微調(diào)方式的測(cè)試結(jié)果也表明，這種微調(diào)增強(qiáng)了DT-Mem適應(yīng)未知情景的能力。

需要說明的是，下表中用來測(cè)試的游戲?qū)τ贛DT來說是已知的，因此MDT的表現(xiàn)在這一輪當(dāng)中不作為衡量依據(jù)。

除了玩游戲，團(tuán)隊(duì)還使用了Meta-World ML45基準(zhǔn)對(duì)DT-Mem進(jìn)行了測(cè)試。

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

結(jié)果顯示，未經(jīng)微調(diào)的模型當(dāng)中，DT-Mem成績(jī)比HDT高出8個(gè)百分點(diǎn)。

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

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