雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制

機器之心發表於2020-03-30

在 Transformer 完全採用注意力機制之後,注意力機制有有了哪些改變?雜湊演算法、Head 之間的資訊交流都需要考慮,視訊記憶體佔用、表徵能力都不能忽視。

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制

注意力機制是非常優美而神奇的機制,在神經網路「資訊過載」的今天,讓 NN 學會只關注特定的部分,無疑會大幅度提升任務的效果與效率。藉助注意力機制神經機器翻譯、預訓練語言模型等任務獲得了前所未有的提升。

但與此同時,注意力機制也面臨著重重問題,首先就是引數量太大,這有點類似於全連線層,注意力機制需要考慮兩兩之間的所有連線。我們可以看到,完全用注意力機制的模型,引數量輕輕鬆鬆破個億,而卷積這類引數共享的運算,引數量一般也就幾百萬。

如果只是引數量大也就算了,能完整利用這些引數,正是說明模型的表徵能力非常強。但問題在於,Transformer 採用的 Multi-head Attention 並沒有充分利用引數的表徵能力。舉個例子,Transformer 中每一個注意力 Head 都是相互獨立的,它們之間沒有資訊交流,因此谷歌最近提出的 Talking-Head 就旨在解決這個問題。

本文從原 Multi-head Attention 出發,探索 Reformer 如何用雜湊演算法大量降低視訊記憶體需求,探索 Talking-Head 如何強化全注意力機制的表徵能力

多頭注意力:開始的地方

Transformer 因在大型預訓練語言模型中的優秀效能而被世人所熟知。這一類模型已廣泛應用於多種預訓練語言模型中,如 BERT、GPT-2 等。當然,在廣泛應用之餘,人們也在思考 Transformer 存在的缺陷,並進行彌補。
  • 論文:Attention Is All You Need

  • 論文連結:https://arxiv.org/abs/1706.03762 (https://arxiv.org/abs/1706.03762)

提到 Transformer,人們總會想到這類模型最突出的幾個特徵:點乘注意力、多頭注意力、殘差連線、Mask 機制等。在最初介紹 Transformer 的論文中,Waswani 等引入了這些優秀的特性,使 Transformer 獨樹一幟。

優秀的特性,問題的根源

在 Transformer 中,首先引入的是注意力點乘機制。它將輸入分為查詢(Query)、鍵(Key)和值(Value),通過以下公式計算。其相當於根據序列元素之間的相似性,確定每一個元素都應該關注哪些資訊。

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制

而另一方面,僅計算一次注意力,不足以捕捉所有的文字結構特徵。因此,Transformer 的提出者想到類似於 CNN 中的 multi-channel 的方法:讓點乘注意力機制「重複」多次。這樣一來,Transformer 的模型複雜度顯著上升,捕捉了更多的語義和語言特徵,這種方法便是我們熟悉的多頭注意力。

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制

圖 2:(左)點乘注意力;(右)多頭注意力,有多個注意力層並行。

多頭注意力的公式如下所示,每一個 Head 都是一個注意力層。為了提高差異性,每一個 Head 的輸入張量都會先經過它特有的線性變換,相當於希望該 Head 注意特定的某個方面。

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制

儘管 Transformer 因為這些特性而產生了很好的效能,但其固有的缺陷也隨之而生。特別需要注意的就是多頭注意力機制了。在設計之初,為了讓不同的注意力頭能夠捕捉語言模型中的不同特徵,不論是 BERT 還是 GPT-2 的設計者都將增加了注意力頭的數量。

但是,由於 Query、Key 和 Value 在訓練中會進行多次線性變換,且每個頭的注意力單獨計算,因此產生了大量的冗餘引數,也需要極大地視訊記憶體。

怎樣才能降低計算複雜度呢?研究者們在這個問題上思考了很多方法,本文中可以看到採用雜湊演算法、引數壓縮等的方法。這些方法有的是強化多頭注意力機制。有的則是利用雜湊演算法的高效性,替換掉原本的架構,使模型效能更加優秀。

雜湊改造下的注意力機制

Transformer 因為 Multi-Head Attention 而成功,也因為它而受到重重限制。在 2017 年 Transformer 提出以後,18 年到 19 年已經有很多研究在改進這種注意力機制,這裡介紹的是一種「雜湊大法」,提出該方法的 Reformer 已經被接收為 ICLR 2020 的 Oral 論文。

ICLR 2020 程式主席評論道:「Reformer 這篇論文提出的新型注意力機制有效減少序列長度的複雜度,同時新機制也減少了儲存需求。實驗證明它們非常有效,該論文可能對研究社群會產生顯著影響。儘管有一些次要的觀點,但所有評審者都一致同意接收該論文。」
  • 論文:Reformer: The Efficient Transformer

  • 論文地址:https://openreview.net/pdf?id=rkgNKkHtvB

注意力機制到底哪不對?

我們先回到注意力機制,其中非常重要的運算是 Query 與 Key 這兩個張量之間的矩陣乘法,其代表著餘弦相似性。具體而言,如果 Query 表示翻譯領域的中文序列,Key 表示英文序列,那麼矩陣乘法則表示在翻譯成某個中文詞時,需要注意哪些英文詞。

但問題在於,Query 與 Key 的張量維度都是 [batch size, length, d_model],乘完後的維度是 [batch size, length, length]。要是序列長度 length 不長也就罷了,但如果上下文跨度很廣,需要很長的序列,也就是說 length 遠遠大於隱藏層大小 d_model。那可就麻煩了,研究者表示要是序列長度達到 64K,注意力機制計算一個樣本也需要 16GB 的視訊記憶體(Float32)。

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制

你可能會說,我不一次性計算 Query 中的所有序列,而每次只算一步,例如一箇中文詞,計算它與所有英文詞之間的相似性。這樣不就沒問題了麼?這是沒問題的,但並沒有實質上的改變。

雜湊來幫忙

現在,假定序列長度還是 64K,對於 Query 序列中的每一個 q,正常的注意力機制需要計算它和 Key 序列 64K 個元素之間的相似度,並通過 Softmax 將相似度歸一化為概率。反正都是要計算概率,且一般只有概率最高的一些元素真正對 q 有很大的貢獻,那麼為什麼不直接找出這些元素?

之前超大的矩陣乘法會計算 Query 序列所有元素與 Key 序列所有元素之間的相似度,現在如果不通過矩陣乘法,只找每個 Query 序列「最相近」的 32 個或 64 個元素,那麼視訊記憶體與計算豈不是成千倍地減少?

這就是 Reformer 最核心的思想,完成查詢「最相近」元素的演算法即區域性敏感雜湊演算法(Locality sensitive hashing)。

簡單而言,區域性敏感雜湊演算法(LSH)在輸入資料彼此類似時,它們有很大概率對映後的雜湊是一樣的;而當輸入資料彼此不同,它們對映後的雜湊值相等概率極小。

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制

LSH 演算法根據區域性敏感雜湊函式族將類似的資料對映到相同的桶(Bucket)。

LSH 注意力機制

如下圖 a 所示,傳統注意力機制需要計算 q 與 k 之間的所有相關性,但實際上真正起作用的注意力分配是非常稀疏的。對於圖 b 來說,Query 和 Key 會先通過雜湊桶排序,相似的 q 與 k 會落在相同的桶內,因此直接在桶內執行注意力計算就能逼近完整的注意力機制

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制

然而圖 b 只是理想情況,真正用於注意力機制還有很多困難要克服,比如說,一個桶內包含很多 q 而沒有一個 k。為此,研究者對 b 執行了一系列轉換,並使排序後的注意力權重,來自相同桶的 q 和 v 都集中在對角線上,就如同圖 c 一樣。注意 c 圖中 Q=K,表示 Transformer 中的自注意力機制

排好序後,d 圖就可以分成不同的批量,並完成平行計算。這樣不僅計算量與引數量大大減少,還能加速訓練。

最後,整個 LSH 注意力機制就如上左圖所示。對於自注意力機制,先根據 LSH 分成不同的桶,然後對桶排序並分割為不同的批量,這些批量能夠並行處理。最後只要在相同的桶內執行注意力機制就行了,整個過程會節省大量儲存與計算資源。

Talking-Heads Attention

近日,來自 Google 的研究團隊提出一種「交談注意力機制」(Talking-Heads Attention),在 softmax 操作前後引入對多頭注意力之間的線性對映,以此增加多個注意力機制間的資訊交流。這樣的操作雖然增加了模型的計算複雜度,卻能夠在多項語言處理問題上取得更好的效果。
  • 論文:Talking-Heads Attention

  • 論文地址:https://arxiv.org/abs/2003.02436

einsum 表示法

作者在論文的虛擬碼中使用大寫字母表示張量,而小寫字母表示其對應維度大小。同時作者在張量的計算中使用了 einsum 表示法,也就是愛因斯坦求和約定。它在 numpy、tensorflow、pytorch 等 Python 擴充套件庫中均有實現。使用 einsum 表示法能夠僅使用一個函式,就可以優雅地實現如點積、外積、轉置、矩陣-向量乘法、矩陣-矩陣乘法等運算。

einsum 解釋:https://rockt.github.io/2018/04/30/einsum

限於本文篇幅,有關 einsum 表示法的概念,感興趣的小夥伴可參考上面這篇部落格。這裡舉個例子,兩個矩陣的乘法運算使用 einsum 表示法可寫為:

Y[a, c] = einsum(X[a, b], W[b, c])

於是,前面介紹的多頭注意力機制使用 einsum 表示法可改寫為如下形式:

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制

 同時,einsum 表示法還支援大於兩個矩陣作為輸入的運算。於是,以上虛擬碼可進一步精簡為如下極簡模式:

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制

交談注意力機制

不同於多頭注意力機制,交談注意力機制在 softmax 運算的前後增加了兩個可學習的線性對映 P_l 與 P_w,這使得多個注意力 logit 和注意力 weight 彼此之間能夠相互傳遞資訊。

於是,交談注意力機制出現了三個不同的維度:h_k、h 和 h_v,其中 h_k 表示 Query 和 Key 的注意力頭數量,h 表示 logit 和 weight 的注意力頭數量,h_v 則表示值的注意力頭數量。其對應虛擬碼表示如下,註釋中標出了每個 einsum 運算所對應的計算量。

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制

Talking-Head 是 Multi-Head 的延伸

當假設注意力機制的輸入與輸出維度相同時,可得到計算多頭注意力機制所需進行的標量乘法運算數目為:

h·(dk +dv)·(n·dX +m·dM +n·m)

而交談注意力機制所需的標量乘法運算數目為:

(dk ·hk +dv ·hv)·(n·dX +m·dM +n·m)+n·m·h·(hk +hv)

可以看到上式中第一項與多頭注意力機制的計算量類似,而第二項是由交談注意力中的線性對映導致的。假如 h<d_k 且 h<d_v,那麼交談注意力機制的線性對映計算量,將分別小於對應多頭注意力機制中的 n·m·dk ·hk 與 n·m·dv ·hv。

作者指出,由於交談注意力機制引入了額外的小尺寸張量運算,在實際執行過程中很可能導致速度比 Transformer 更慢。

最後,多頭注意力機制與交談注意力機制都可看作一種「通用雙線性多頭注意力機制」(GBMA, i.e. general bilinear multihead attention)的特殊形式。我們可以從以下虛擬碼看出,GBMA 具有兩個三維的引數張量

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制
從以下虛擬碼不難看出,多頭注意力機制在數學上等效於,使用兩個因子的乘積分別表示 GBMA 中的各引數張量

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制

而交談注意力機制在數學上等效於,使用三個因子的乘積分別表示 GBMA 中的各引數張量

雜湊演算法、愛因斯坦求和約定,這是2020年的注意力機制


這裡的 GBMA 僅作為理論研究探討,由於其計算量較大可能並不具備實用性。

「注意力」這個想法真的非常優雅,所以 2020 年剩下的是,我們如何才能更高效地實現「注意力」?


相關文章