Bengio等人新作：注意力可被視為RNN，新模型媲美Transformer，但超級(jí)省內(nèi)存
序列建模的進(jìn)展具有極大的影響力，因?yàn)樗鼈冊(cè)趶V泛的應(yīng)用中發(fā)揮著重要作用，包括強(qiáng)化學(xué)習(xí)（例如，機(jī)器人和自動(dòng)駕駛）、時(shí)間序列分類（例如，金融欺詐檢測(cè)和醫(yī)學(xué)診斷）等。
在過去的幾年里，Transformer 的出現(xiàn)標(biāo)志著序列建模中的一個(gè)重大突破，這主要得益于 Transformer 提供了一種能夠利用 GPU 并行處理的高性能架構(gòu)。
然而，Transformer 在推理時(shí)計(jì)算開銷很大，主要在于內(nèi)存和計(jì)算需求呈二次擴(kuò)展，從而限制了其在低資源環(huán)境中的應(yīng)用（例如，移動(dòng)和嵌入式設(shè)備）。盡管可以采用 KV 緩存等技術(shù)提高推理效率，但 Transformer 對(duì)于低資源領(lǐng)域來說仍然非常昂貴，原因在于：（1）隨 token 數(shù)量線性增加的內(nèi)存，以及（2）緩存所有先前的 token 到模型中。在具有長(zhǎng)上下文（即大量 token）的環(huán)境中，這一問題對(duì) Transformer 推理的影響更大。
為了解決這個(gè)問題，加拿大皇家銀行 AI 研究所 Borealis AI、蒙特利爾大學(xué)的研究者在論文《Attention as an RNN 》中給出了解決方案。值得一提的是，我們發(fā)現(xiàn)圖靈獎(jiǎng)得主 Yoshua Bengio 出現(xiàn)在作者一欄里。
- 論文地址：https://arxiv.org/pdf/2405.13956
- 論文標(biāo)題：Attention as an RNN?
具體而言，研究者首先檢查了 Transformer 中的注意力機(jī)制，這是導(dǎo)致 Transformer 計(jì)算復(fù)雜度呈二次增長(zhǎng)的組件。該研究表明注意力機(jī)制可以被視為一種特殊的循環(huán)神經(jīng)網(wǎng)絡(luò)（RNN），具有高效計(jì)算的多對(duì)一（many-to-one）RNN 輸出的能力。利用注意力的 RNN 公式，該研究展示了流行的基于注意力的模型（例如 Transformer 和 Perceiver）可以被視為 RNN 變體。
然而，與 LSTM、GRU 等傳統(tǒng) RNN 不同，Transformer 和 Perceiver 等流行的注意力模型雖然可以被視為 RNN 變體。但遺憾的是，它們無法高效地使用新 token 進(jìn)行更新。
為了解決這個(gè)問題，該研究引入了一種基于并行前綴掃描（prefix scan）算法的新的注意力公式，該公式能夠高效地計(jì)算注意力的多對(duì)多（many-to-many）RNN 輸出，從而實(shí)現(xiàn)高效的更新。
在此新注意力公式的基礎(chǔ)上，該研究提出了 Aaren（[A] ttention [a] s a [re] current neural [n] etwork），這是一種計(jì)算效率很高的模塊，不僅可以像 Transformer 一樣并行訓(xùn)練，還可以像 RNN 一樣高效更新。
實(shí)驗(yàn)結(jié)果表明，Aaren 在 38 個(gè)數(shù)據(jù)集上的表現(xiàn)與 Transformer 相當(dāng)，這些數(shù)據(jù)集涵蓋了四種常見的序列數(shù)據(jù)設(shè)置：強(qiáng)化學(xué)習(xí)、事件預(yù)測(cè)、時(shí)間序列分類和時(shí)間序列預(yù)測(cè)任務(wù)，同時(shí)在時(shí)間和內(nèi)存方面更加高效。
方法介紹
為了解決上述問題，作者提出了一種基于注意力的高效模塊，它能夠利用 GPU 并行性，同時(shí)又能高效更新。
首先，作者在第 3.1 節(jié)中表明，注意力可被視為一種 RNN，具有高效計(jì)算多對(duì)一 RNN（圖 1a）輸出的特殊能力。利用注意力的 RNN 形式，作者進(jìn)一步說明，基于注意力的流行模型，如 Transformer（圖 1b）和 Perceiver（圖 1c），可以被視為 RNN。然而，與傳統(tǒng)的 RNN 不同的是，這些模型無法根據(jù)新 token 有效地更新自身，從而限制了它們?cè)跀?shù)據(jù)以流的形式到達(dá)的序列問題中的潛力。
為了解決這個(gè)問題，作者在第 3.2 節(jié)中介紹了一種基于并行前綴掃描算法的多對(duì)多 RNN 計(jì)算注意力的高效方法。在此基礎(chǔ)上，作者在第 3.3 節(jié)中介紹了 Aaren—— 一個(gè)計(jì)算效率高的模塊，它不僅可以并行訓(xùn)練（就像 Transformer），還可以在推理時(shí)用新 token 高效更新，推理只需要恒定的內(nèi)存（就像傳統(tǒng) RNN）。
將注意力視為一個(gè)多對(duì)一 RNN
查詢向量 q 的注意力可被視為一個(gè)函數(shù)，它通過 N 個(gè)上下文 token x_1:N 的鍵和值
將其映射到單一輸出 o_N = Attention (q, k_1:N , v_1:N ) 。給定 s_i = dot (q，k_i)，輸出 o_N 可表述為：
，分母為
。將注意力視為 RNN，可以在 k = 1，...，...... 時(shí)，以滾動(dòng)求和的方式迭代計(jì)算
和
。然而，在實(shí)踐中，這種實(shí)現(xiàn)方式并不穩(wěn)定，會(huì)因有限的精度表示和可能非常小或非常大的指數(shù)（即 exp (s)）而遇到數(shù)值問題。為了緩解這一問題，作者用累積最大值項(xiàng)
來重寫遞推公式，計(jì)算
和
。值得注意的是，最終結(jié)果是相同的
，m_k 的循環(huán)計(jì)算如下：
計(jì)算注意力的方法：通過將注意力視為一個(gè) RNN，可以看到計(jì)算注意力的不同方法：在 O (1) 內(nèi)存中逐個(gè) token 循環(huán)計(jì)算（即順序計(jì)算）；或以傳統(tǒng)方式計(jì)算（即并行計(jì)算），需要線性 O (N) 內(nèi)存。由于注意力可以被看作是一個(gè) RNN，因此計(jì)算注意力的傳統(tǒng)方法也可以被看作是計(jì)算注意力多對(duì)一 RNN 輸出的高效方法，即 RNN 的輸出以多個(gè)上下文 token 為輸入，但在 RNN 結(jié)束時(shí)只輸出一個(gè) token（見圖 1a）。最后，也可以將注意力計(jì)算為一個(gè)逐塊處理 token 的 RNN，而不是完全按順序或完全并行計(jì)算，這需要 O (b) 內(nèi)存，其中 b 是塊的大小。
將現(xiàn)有的注意力模型視為 RNN。通過將注意力視為 RNN，現(xiàn)有的基于注意力的模型也可以被視為 RNN 的變體。例如，Transformer 的自注意力是 RNN（圖 1b），上下文 token 是其初始隱藏狀態(tài)。Perceiver 的交叉注意力是 RNN（圖 1c），其初始隱藏狀態(tài)是與上下文相關(guān)的潛變量。通過利用其注意力機(jī)制的 RNN 形式，這些現(xiàn)有模型可以高效地計(jì)算其輸出存儲(chǔ)。
然而，當(dāng)將現(xiàn)有的基于注意力的模型（如 Transformers）視為 RNN 時(shí)，這些模型又缺乏傳統(tǒng) RNN（如 LSTM 和 GRU）中常見的重要屬性。
值得注意的是，LSTM 和 GRU 能夠僅在 O (1) 常量?jī)?nèi)存和計(jì)算中使用新 token 有效地更新自身，相比之下， Transformer 的 RNN 視圖（見圖 1b）會(huì)通過將一個(gè)新的 token 作為初始狀態(tài)添加一個(gè)新的 RNN 來處理新 token。這個(gè)新的 RNN 處理所有先前的 token，需要 O (N) 的線性計(jì)算量。
在 Perceiver 中，由于其架構(gòu)的原因，潛變量（圖 1c 中的 L_i）是依賴于輸入的，這意味著它們的值在接收新 token 時(shí)會(huì)發(fā)生變化。由于其 RNN 的初始隱藏狀態(tài)（即潛變量）發(fā)生變化，Perceiver 因此需要從頭開始重新計(jì)算其 RNN，需要 O (NL) 的線性計(jì)算量，其中 N 是 token 的數(shù)量，L 是潛變量的數(shù)量。
將注意力視為一個(gè)多對(duì)多 RNN
針對(duì)這些局限性，作者建議開發(fā)一種基于注意力的模型，利用 RNN 公式的能力來執(zhí)行高效更新。為此，作者首先引入了一種高效的并行化方法，將注意力作為多對(duì)多 RNN 計(jì)算，即并行計(jì)算
的方法。為此，作者利用并行前綴掃描算法（見算法 1），這是一種通過關(guān)聯(lián)算子 ⊕ 從 N 個(gè)連續(xù)數(shù)據(jù)點(diǎn)計(jì)算 N 個(gè)前綴的并行計(jì)算方法。該算法可高效計(jì)算
回顧
，其中
，
為了高效計(jì)算
，可以通過并行掃描算法計(jì)算
和
，然后結(jié)合 a_k 和 c_k 計(jì)算
。
為此，作者提出了以下關(guān)聯(lián)算子⊕，該算子作用于形式為（m_A、u_A、w_A）的三元組，其中 A 是一組索引，
，
，
。并行掃描算法的輸入為
。該算法遞歸應(yīng)用算子 ⊕，其工作原理如下：
，其中，
，
。
在完成遞歸應(yīng)用算子后，算法輸出
。也被稱作
。結(jié)合輸出元組的最后兩個(gè)值，檢索
從而產(chǎn)生一種高效的并行方法，將注意力計(jì)算為多對(duì)多 RNN（圖 3）。
Aaren：[A] ttention [a] s a [re] current neural [n] etwork
Aaren 的接口與 Transformer 相同，即將 N 個(gè)輸入映射到 N 個(gè)輸出，而第 i 個(gè)輸出是第 1 到第 i 個(gè)輸入的聚合。此外，Aaren 還自然可堆疊，并且能夠計(jì)算每個(gè)序列 token 的單獨(dú)損失項(xiàng)。然而，與使用因果自注意力的 Transformers 不同，Aaren 使用上述計(jì)算注意力的方法作為多對(duì)多 RNN，使其更加高效。Aaren 形式如下：
與 Transformer 不同，在 Transformer 中查詢是輸入到注意力的 token 之一，而在 Aaren 中，查詢 token q 是在訓(xùn)練過程中通過反向傳播學(xué)習(xí)得到的。
下圖展示了一個(gè)堆疊 Aaren 模型的例子，該模型的輸入上下文 token 為 x_1:3，輸出為 y_1:3。值得注意的是，由于 Aaren 利用了 RNN 形式的注意力機(jī)制，堆疊 Aarens 也相當(dāng)于堆疊 RNN。因此，Aarens 也能夠高效地用新 token 進(jìn)行更新，即 y_k 的迭代計(jì)算僅需要常量計(jì)算，因?yàn)樗鼉H依賴于 h_k-1 和 x_k。
基于 Transformer 的模型需要線性內(nèi)存（使用 KV 緩存時(shí)）并且需要存儲(chǔ)所有先前的 token ，包括中間 Transformer 層中的那些，但基于 Aaren 的模型只需要常量?jī)?nèi)存，并且不需要存儲(chǔ)所有先前的 token ，這使得 Aarens 在計(jì)算效率上顯著優(yōu)于 Transformer。
實(shí)驗(yàn)
實(shí)驗(yàn)部分的目標(biāo)是比較 Aaren 和 Transformer 在性能和所需資源（時(shí)間和內(nèi)存）方面的表現(xiàn)。為了進(jìn)行全面比較，作者在四個(gè)問題上進(jìn)行了評(píng)估：強(qiáng)化學(xué)習(xí)、事件預(yù)測(cè)、時(shí)間序列預(yù)測(cè)和時(shí)間序列分類。
強(qiáng)化學(xué)習(xí)
作者首先比較了 Aaren 和 Transformer 在強(qiáng)化學(xué)習(xí)方面的表現(xiàn)。強(qiáng)化學(xué)習(xí)在機(jī)器人、推薦引擎和交通控制等交互式環(huán)境中很受歡迎。
表 1 中的結(jié)果表明，在所有 12 個(gè)數(shù)據(jù)集和 4 種環(huán)境中，Aaren 與 Transformer 的性能都不相上下。不過，與 Transformer 不同的是，Aaren 也是一種 RNN，因此能夠在持續(xù)計(jì)算中高效處理新的環(huán)境交互，從而更適合強(qiáng)化學(xué)習(xí)。
事件預(yù)測(cè)
接下來，作者比較了 Aaren 和 Transformer 在事件預(yù)測(cè)方面的表現(xiàn)。事件預(yù)測(cè)在許多現(xiàn)實(shí)環(huán)境中都很流行，例如金融（如交易）、醫(yī)療保?。ㄈ缁颊哂^察）和電子商務(wù)（如購(gòu)買）。
表 2 中的結(jié)果顯示，Aaren 在所有數(shù)據(jù)集上的表現(xiàn)都與 Transformer 相當(dāng)。Aaren 能夠高效處理新輸入，這在事件預(yù)測(cè)環(huán)境中尤為有用，因?yàn)樵谶@種環(huán)境中，事件會(huì)以不規(guī)則流的形式出現(xiàn)。
時(shí)間序列預(yù)測(cè)
然后，作者比較了 Aaren 和 Transformer 在時(shí)間序列預(yù)測(cè)方面的表現(xiàn)。時(shí)間序列預(yù)測(cè)模型通常用在與氣候（如天氣）、能源（如供需）和經(jīng)濟(jì)（如股票價(jià)格）相關(guān)的領(lǐng)域。
表 3 中的結(jié)果顯示，在所有數(shù)據(jù)集上，Aaren 與 Transformer 的性能相當(dāng)。不過，與 Transformer 不同的是，Aaren 能高效處理時(shí)間序列數(shù)據(jù)，因此更適合與時(shí)間序列相關(guān)的領(lǐng)域。
時(shí)間序列分類
接下來，作者比較了 Aaren 和 Transformer 在時(shí)間序列分類方面的表現(xiàn)。時(shí)間序列分類在許多重要的應(yīng)用中很常見，例如模式識(shí)別（如心電圖）、異常檢測(cè)（如銀行欺詐）或故障預(yù)測(cè)（如電網(wǎng)波動(dòng)）。
從表 4 中可以看出，在所有數(shù)據(jù)集上，Aaren 與 Transformer 的表現(xiàn)不相上下。
分析
最后，作者比較了 Aaren 和 Transformer 所需的資源。
內(nèi)存復(fù)雜性：在圖 5（左）中，作者比較了 Aaren 和 Transformer（使用 KV 緩存）在推理時(shí)的內(nèi)存使用情況?？梢钥吹剑殡S KV 緩存技術(shù)的使用，Transformer 的內(nèi)存使用量呈線性增長(zhǎng)。相比之下，Aaren 只使用恒定的內(nèi)存，無論 token 數(shù)量如何增長(zhǎng)，因此它的效率要高得多。
時(shí)間復(fù)雜度：在圖 5（右圖）中，作者比較了 Aaren 和 Transformer（使用 KV 緩存）按順序處理一串 token 所需的累計(jì)時(shí)間。對(duì)于 Transformer，累計(jì)計(jì)算量是 token 數(shù)的二次方，即 O (1 + 2 + ... + N) = O (N^2 )。相比之下，Aaren 的累計(jì)計(jì)算量是線性的。在圖中，可以看到模型所需的累計(jì)時(shí)間也是類似的結(jié)果。具體來說，Transformer 所需的累計(jì)時(shí)間呈二次增長(zhǎng)，而 Aaren 所需的累計(jì)時(shí)間呈線性增長(zhǎng)。
參數(shù)數(shù)量：由于要學(xué)習(xí)初始隱藏狀態(tài) q，Aaren 模塊需要的參數(shù)略多于 Transformer 模塊。不過，由于 q 只是一個(gè)向量，因此差別不大。通過在同類模型中進(jìn)行實(shí)證測(cè)量，作者發(fā)現(xiàn) Transformer 使用了 3, 152, 384 個(gè)參數(shù)。相比之下，等效的 Aaren 使用了 3, 152, 896 個(gè)參數(shù)，參數(shù)增加量?jī)H為 0.016%—— 對(duì)于內(nèi)存和時(shí)間復(fù)雜性的顯著差異來說，這只是微不足道的代價(jià)。
本文轉(zhuǎn)自 機(jī)器之心 ，作者：機(jī)器之心
