文 | 小軼

今天給大家介紹一篇谷歌的最新工作，解決的是 Transformer 的長文本處理問題。在原生 Transformer 中，attention 的復雜度是輸入序列長度的平方級別，因此限制了它處理長文本的能力。簡單來說，本文提出的解決方案就是把 Transformer當做 RNN 中的循環單元來用。

和傳統 RNN 的區別只在于：傳統 RNN encoder 每個循環單元負責編碼一個 token，而本文中每個循環單元負責編碼一段長度為 的文本片段，且每個循環單元都由構造相同的 Transformer Block 來實現。如此一來，每個片段在編碼時，都能用類似 RNN 的方式，同時考慮之前文本中的信息了。

想法很簡單，但具體實現起來還是有一些難點。接下來，我們展開介紹一下本文所提出的 Block-Recurrent Transformer。

論文標題：
BLOCK-RECURRENT TRANSFORMERS

論文鏈接：
https://arxiv.org/pdf/2203.07852.pdf

滑動注意力機制

先來看一下每個 block 的 attention 范圍。本文采用的是一種滑動窗口注意力機制，一種專門針對長文檔場景的技術。由于文本過長，讓每個 token 都 attend 到整個文本中的所有 token 難以實現。在滑動窗口注意力機制中：每個 token 只需要 attend 到它的前 個 token。在本文中，滑動窗口長度 與每個循環單元所需處理的文本長度 相等，即 。

上圖示例中，假設窗口長度為 8；相應地，輸入文本也被分為長度為 8 的片段，交由 Transformer blocks 分別處理。圖中淺藍色區域表示了 attention 范圍。

圖中兩個黑框分別對應了兩個 Transformer block 。8個紅色標記點，代表右下角那個 block 所需要處理的 8 個 token。可以看到，每個 block 的 attention 矩陣大小為 。因此，對于長度為 N 的輸入來說，整個模型的 attention 復雜度為 O(N)。

循環單元

接下來，我們就往每個 Transformer block 內部看看，究竟是如何實現循環的。

▲傳統 RNN 結構

類似傳統 RNN，每個循環單元：

• 輸入是input embeddings 和 current state

• 輸出是 output embeddings 和 next state

所以，我們這里所需要理解的兩個核心問題也就是：在 Block-Recurrent Transformer 中，這兩個輸出分別是如何得到的？

垂直方向：如何得到 output embeddings？

下圖展示了得到 output embeddings 的過程。

▲垂直方向：如何得到 output embeddings

和傳統的 Transformer layer 非常相像，差別只集中在紅框標識出來的部分。在這一部分中，為了融合上一個循環單元給的 current state 信息，他們將 input embeddings 和 current state vectors 做了一個 cross attention。另一方面，input embeddings 自身也會過一個 self-attention 層。這兩部分拼接后，通過線性層融合在了一起。

水平方向：如何得到 next state？

下圖展示了得到 next state 的過程。

▲水平方向：循環單元之間如何傳遞 state

與傳統 Transformer 不同的地方，用紅色和粉色框標識了。紅色部分，同樣是用 cross attention 將 input embeddings 和 current state vectors 融合。粉色部分則是用兩個 gate 替代了原本 Transformer 中的殘差層。這兩個 gate 的作用與 LSTM 中的遺忘門類似，控制了對前一個 state 信息的保留程度。

垂直方向如何多層疊加？

最后還有一個問題。我們都知道，傳統 Transformer Encoder 通常是由多個 Transformer Layer 疊加起來的。也就是下圖中那個 的意義。那么，在 Block-Recurrent Transformer 中，如何實現垂直方向上的多層疊加呢？

▲傳統 Transformer Encoder

文中討論了兩種方式，Single Recurrent Layer 和 Feedback。

▲Single Recurrent Layer

Single Recurrent Layer （SRL） 的實現比較簡單。我簡單花了張示意圖，大致如上圖所示。垂直方向上疊加的多個層：大多數都是普通的 Transformer Layer；只有其中的一層，在水平方向上接收了 current state，做了循環操作。這種方式的運算復雜度也比較低，只相當于在普通的 Transformer 基礎上多加了一層 layer 的運算量。也就是說，如果垂直疊加了 12 層，相當于普通 Transformer 疊加 13 層的運算量。

▲Feedback

Feedback 在 SRL 的基礎上，current state 還會廣播給其他 Transformer Layer。這些層會用 cross attention 的方式，將 current state 的信息融合。實驗中，Feedback 比 SRL 性能有小幅提升，不過它的模型參數更多，訓練時長也要陡增 35~40%。

實驗

實驗在三個長文本數據集上進行，分別是 PG19，arxiv 和 Github。評測任務是自回歸語言建模，指標為 perplexity。結果如下圖所示。

其中，黃色高亮的是本文所提出方法的兩個變種，獲得了 SOTA 的效果。

紅色框出的是三個比較重要的 baseline。其中，上面兩個 baseline 是此前經典的長文檔處理模型 Transformer-XL 的兩個變種。可以看到本文方法的性能要比他們好不少。

最后一行的 Memorizing Transformer 同樣是谷歌的工作，剛剛被 ICLR'2022 錄用。其基本思想是：編碼長文本時，模型一邊往下讀，一邊把之前見過的所有 token 保存在一個數據庫中；在讀當前片段時，會用 kNN 的方式找到數據庫中相似的內容，然后和當前內容同時交互編碼。

可以看到，這個模型的效果其實和本文方法相差不大，但復雜度要高很多，運算時延也要長[1]。雖然...但是，本文并沒有把 Memorizing Transformer 的 step time 明確寫在表格中。個人感覺有些不妥。

小結

本文的想法其實很簡單：把 Transformer 作為 RNN 的循環單元，解決長文本問題。我相信想到過類似 idea 的應該早有人在。我確實也看到了類似的 previous works，不過它們的模型復雜度和性能效果都遜于本文。

就本文來說，只是擁有一個 idea 肯定是不夠的，還要解決很多問題，包括：

• 相鄰的 block 之間如何以適配 Transformer 的方式傳遞信息

• 模型設計的時候還要同時考慮到將運算復雜度的降到最低，能并行運算的絕不搞串行

• 還有最后工程實現上的一些問題。比如說，模型訓練的時候是否會像傳統 RNN 一樣遇到梯度消失的問題？如果有，該如何解決？我在本篇推送中，沒有涵蓋這方面的討論。原文確實提了一些方法來提高模型訓練的穩定性。

從一個宏觀的 idea 到真正落實，還是有很長距離的。所以還是不能輕易地說一篇論文的 idea “too simple”。

往期回顧

《Longformer：超越RoBERTa，為長文檔而生的預訓練模型》

《告別自注意力，谷歌為Transformer打造新內核Synthesizer》

《Google綜述：細數Transformer模型的17大高效變種》

萌屋作者：小軼

是小軼，不是小秩！更不要叫小鐵！高冷的形象是需要大家共同維護的！作為成熟的大人，正在勤儉節約、兢兢業業，為成為一名合格的（但是仍然發量充足的）PhD而努力著。日常沉迷對話系統。說不定，正在和你對話的，并不是不是真正的小軼哦（！？）

“高冷？那是站在冰箱頂端的意思啦。” ?——白鹡鸰

作品推薦：

寫了一篇關于 NLP 綜述的綜述！

全球44家機構，55位大佬，歷時兩年，打造最強NLG評測基準！

谷歌重磅：可以優化自己的優化器！手動調參或將成為歷史！？

ACL20 Best Paper揭曉！NLP模型評價體系或將迎來重大轉折

后臺回復關鍵詞【入群】

加入賣萌屋NLP、CV與搜推廣求職討論群

后臺回復關鍵詞【頂會】

獲取ACL、CIKM等各大頂會論文集！

?

[1] Memorizing Transformers https://arxiv.org/abs/2203.08913

總結

以上是生活随笔為你收集整理的谷歌提出 RNN 版 Transformer，或为长文本建模的当前最优解的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。