視覺Transformer中ReLU替代softmax

點擊上方“小白學視覺”，選擇加"星標"或“置頂”

重磅干貨，第一時間送達

Transformer 架構(gòu)已經(jīng)在現(xiàn)代機器學習領(lǐng)域得到了廣泛的應(yīng)用。注意力是 transformer 的一大核心組件，其中包含了一個 softmax，作用是產(chǎn)生 token 的一個概率分布。softmax 有較高的成本，因為其會執(zhí)行指數(shù)計算和對序列長度求和，這會使得并行化難以執(zhí)行。

Google DeepMind 想到了一個新思路：用某種不一定會輸出概率分布的新方法替代 softmax 運算。他們還觀察到：在用于視覺 Transformer 時，使用 ReLU 除以序列長度的注意力可以接近或匹敵傳統(tǒng)的 softmax 注意力。

論文：https://arxiv.org/abs/2309.08586

這一結(jié)果為并行化帶來了新方案，因為 ReLU 注意力可以在序列長度維度上并行化，其所需的 gather 運算少于傳統(tǒng)的注意力。

方法

注意力

注意力的作用是通過一個兩步式流程對 d 維的查詢、鍵和值 {q_i, k_i, v_i} 進行變換。

在第一步，通過下式得到注意力權(quán)重 ：

其中 ? 通常是 softmax。

下一步，使用這個注意力權(quán)重來計算輸出 這篇論文探索了使用逐點式計算的方案來替代 ?。

ReLU 注意力

DeepMind 觀察到，對于 1 式中的 ? = softmax，是一個較好的替代方案。他們將使用 的注意力稱為 ReLU 注意力。

已擴展的逐點式注意力

研究者也通過實驗探索了更廣泛的  選擇，其中 α ∈ [0, 1] 且 h ∈ {relu,relu2 , gelu,softplus, identity,relu6,sigmoid}。

序列長度擴展

他們還觀察到，如果使用一個涉及序列長度 L 的項進行擴展，有助于實現(xiàn)高準確度。之前試圖去除 softmax 的研究工作并未使用這種擴展方案。

在目前使用 softmax 注意力設(shè)計的 Transformer 中，有  ，這意味著  盡管這不太可能是一個必要條件，但 能確保在初始化時  的復(fù)雜度是  ，保留此條件可能會減少替換 softmax 時對更改其它超參數(shù)的需求。

在初始化的時候，q 和 k 的元素為 O (1)，因此  也將為 O (1)。ReLU 這樣的激活函數(shù)維持在 O (1)，因此需要因子才能使  的復(fù)雜度為  。

實驗與結(jié)果

主要結(jié)果

圖 1 說明在 ImageNet-21k 訓練方面，ReLU 注意力與 softmax 注意力的擴展趨勢相當。X 軸展示了實驗所需的內(nèi)核計算總時間（小時）。ReLU 注意力的一大優(yōu)勢是能在序列長度維度上實現(xiàn)并行化，其所需的 gather 操作比 softmax 注意力更少。

序列長度擴展的效果

圖 2 對比了序列長度擴展方法與其它多種替代 softmax 的逐點式方案的結(jié)果。具體來說，就是用 relu、relu2、gelu、softplus、identity 等方法替代 softmax。X 軸是 α。Y 軸則是 S/32、S/16 和 S/8 視覺 Transformer 模型的準確度。最佳結(jié)果通常是在 α 接近 1 時得到。由于沒有明確的最佳非線性，所以他們在主要實驗中使用了 ReLU，因為它速度更快。

qk-layernorm 的效果

主要實驗中使用了 qk-layernorm，在這其中查詢和鍵會在計算注意力權(quán)重前被傳遞通過 LayerNorm。DeepMind 表示，默認使用 qk-layernorm 的原因是在擴展模型大小時有必要防止不穩(wěn)定情況發(fā)生。圖 3 展示了移除 qk-layernorm 的影響。這一結(jié)果表明 qk-layernorm 對這些模型的影響不大，但當模型規(guī)模變大時，情況可能會不一樣。

添加門的效果

先前有移除 softmax 的研究采用了添加一個門控單元的做法，但這種方法無法隨序列長度而擴展。具體來說，在門控注意力單元中，會有一個額外的投影產(chǎn)生輸出，該輸出是在輸出投影之前通過逐元素的乘法組合得到的。圖 4 探究了門的存在是否可消除對序列長度擴展的需求?？傮w而言，DeepMind 觀察到，不管有沒有門，通過序列長度擴展都可以得到最佳準確度。也要注意，對于使用 ReLU 的 S/8 模型，這種門控機制會將實驗所需的核心時間增多大約 9.3%。

      

       下載1：OpenCV-Contrib擴展模塊中文版教程
       

      
      

       在「小白學視覺」公眾號后臺回復(fù)：擴展模塊中文教程，即可下載全網(wǎng)第一份OpenCV擴展模塊教程中文版，涵蓋擴展模塊安裝、SFM算法、立體視覺、目標跟蹤、生物視覺、超分辨率處理等二十多章內(nèi)容。
      
      

       

      
      

       下載2：Python視覺實戰(zhàn)項目52講
      
      

       在「小白學視覺」公眾號后臺回復(fù)：Python視覺實戰(zhàn)項目，即可下載包括圖像分割、口罩檢測、車道線檢測、車輛計數(shù)、添加眼線、車牌識別、字符識別、情緒檢測、文本內(nèi)容提取、面部識別等31個視覺實戰(zhàn)項目，助力快速學校計算機視覺。
      
      

       

      
      

       下載3：OpenCV實戰(zhàn)項目20講
      
      

       在「小白學視覺」公眾號后臺回復(fù)：OpenCV實戰(zhàn)項目20講，即可下載含有20個基于OpenCV實現(xiàn)20個實戰(zhàn)項目，實現(xiàn)OpenCV學習進階。
      
      

       

      
交流群

歡迎加入公眾號讀者群一起和同行交流，目前有SLAM、三維視覺、傳感器、自動駕駛、計算攝影、檢測、分割、識別、醫(yī)學影像、GAN、算法競賽等微信群（以后會逐漸細分），請掃描下面微信號加群，備注：”昵稱+學校/公司+研究方向“，例如：”張三 + 上海交大 + 視覺SLAM“。請按照格式備注，否則不予通過。添加成功后會根據(jù)研究方向邀請進入相關(guān)微信群。請勿在群內(nèi)發(fā)送廣告，否則會請出群，謝謝理解~