首頁 > 科技

上海交大、華為海思提出X-volution,發力網路核心基礎架構創新

2021-06-07 15:52:16

機器之心報道

機器之心編輯部

卷積和自注意力各有優勢,但二者的有效結合一直是一大難題。為了取二者之長,上海交大 - 華為海思聯合團隊提出了一種名為 X-volution 的新型運算元。該運算元在效能上的顯著提升、計算的通用性與即插即用的特性為深度學習基礎計算單元庫以及 NPU 計算架構的演進提供了一種新的基礎武器。

眾所周知,卷積操作(convolution)與自注意力操作(self-attention)是深度學習兩大核心的基礎網路計算單元(或稱為模型運算元)。卷積操作通過線性乘子,提取影象局部特徵;自注意力操作通過高階乘子運算,提取影象全域 / 局部的特徵關聯特性。兩種運算元成為深度學習兩大重要網路架構演化——CNN 與 Transformer 的計算基石。兩種運算元在影象特徵提取與語義抽象方面的互補性不言而喻:線性 vs. 高階, 局部 vs. 全局。因此,能否設計一種包含這兩種操作的融合運算元並使其發揮互補優勢,一直是深度學習架構研究者熱衷的研究方向之一。

然而,由於卷積運算與自注意力運算在計算模式上的異構性,這項任務存在巨大的挑戰。目前學界中的一些工作也在努力統一兩者,他們主要從拓撲結構組合角度來粗粒度地結合兩種運算元,例如,發表在 ICCV 2019 上的 AA-Net 採用了一種將卷積中部分通道替換為由 self-attention 來處理,然後將卷積和 self-attention 分別處理的特徵連線來達到聯合兩種運算元的目的,這種做法證明了卷積和 self-attention 結合後確實能在分類、檢測、分割等基礎任務上達到比較可觀的效能收益。

然而,粗粒度的組合(本質上就是兩路計算並聯)會導致其組合後網路形態可用性下降。具體來說,卷積和 self-attention 運算模式存在較大差異,兩者同時存在會導致網路結構不規則,進而影響網路推理效率,並不為目前一些工業界通用的晶片計算架構所友好支援。同時組合後的運算元在算力上也存在巨大的挑戰。

論文連結:https://arxiv.org/pdf/2106.02253.pdf

針對這些挑戰,日前,上海交大 - 華為海思聯合團隊在 arXiv 上發表了「X-volution: On the Unification of Convolution and Self-attention」,首次在計算模式上統一了這兩大基礎運算元,並在推理階段歸併成一個簡單的卷積型運算元:X-volution。

X-volution 兼顧卷積與自注意力操作的互補優勢,並且在現有通用網路計算框架上不需要額外運算元支援,也不增加除卷積外的額外算力或影響網路的規範性 / 可用性(即插即用)。

該工作的突破主要受以下思路的啟發:對全局的 self-attention 進行理論分析後,研究者發現在一定條件下(例如影象 / 特徵圖的鄰接畫素滿足馬爾可夫性質),全局的 self-attention 可以通過局部的 self-attention 配合卷積運算來逼近。

具體來說,本文作者提出了一種新型的 self-attention 機制——PSSA。這種機制分為兩個步驟:首先將輸入的特徵沿指定的多個方向進行迴圈移位(採用索引來實現)得到移位後的特徵,然後將移位後的特徵與原特徵通過元素點積獲得變換後的特徵,再對該特徵在局部區域進行加權求和(可採用卷積來替代),至此獲得經過注意力機制處理後的特徵。通過層次堆疊,可以持續地將局部的上下文關係傳播到全局從而實現全局的 self-attention。

值得注意的是,PSSA 實際上將 self-attention 巧妙地轉化為了一個在簡單變換後的特徵上的標準的卷積操作,這從形式上實現了 self-attention 向卷積的統一。利用此逼近式的 self-attention 機制,作者建立了一個多分枝的模組將卷積和 self-attention 整合在一起,這個模組從功能上實現了兩者的統一。

更重要的是,這個多分枝的結構可以利用結構重參數化的方法進行有條件的合併。多分枝結構可以合併為單個卷積,合併後可以獲得一個原子級的運算元,稱為 X-volution(X-volution 的權重可以看作一個靜態卷積權重,以及一個內容相關動態卷積權重的和)。此運算元同時具備了卷積和 self-attention 的特性,且不會影響網路的規範性 / 可用性。

作者在分類、檢測、分割等主流 SOTA 實驗上取得了顯著的效能提升。

圖 1,運算元詳細結構框圖。受結構重參數化思想啟發,X-volution 被設計為訓練和推理階段結構解耦的形式。它的訓練結構時有兩個主要分支(如中間所示),右分支由級聯的卷積和 BN 構成,可為 X-volution 整合卷積的能力。左邊包括 PSSA,它提供近似的全局自注意力特性。完成訓練後,X-volution 可以有條件地重新參數化為一個卷積操作。在推理階段,X-volution 實際上是一個動態卷積運算元,它的權重包括需要線上計算的 attention 動態參數部分和已經訓練和固化的卷積靜態參數部分。

實驗部分

作者將 X-volution 接入到經典的 ResNet 模型中用於 ImageNet 分類、MS COCO 物體檢測、例項分割等關鍵基礎任務並都取得了不俗的提升。為了排除其他因素干擾,實驗中作者所使用的 self-attention 和 PSSA 都沒有新增位置編碼,並且沒有對資料集進行額外的增廣,沒有使用額外的訓練技巧(如:餘弦退火、標籤平滑等)。

ImageNet 分類實驗

對於 ImageNet 圖片分類實驗,作者分別測試了在 ResNet 中三個不同位置接入 X-volution 的結果。將 X-volution 接入到常規的 ResNet 第五階段瓶頸單元的結果如表 1 所示:在 ResNet-34 與 ResNet-50 中均提升不明顯,這是因為在此階段的特徵圖尺寸已經接近卷積核大小。實驗發現在第三階段效果最為突出,分別取得了 1.2% 與 0.9% 的顯著提升。值得注意的是,作者對於 ResNet 改動較小,但是效能卻依然能有大幅度的提升,這證實了文中所提出的 X-volution 運算元具有良好的效能。

表 1. ImageNet 實驗結果及瓶頸單元詳細結構

MS COCO 物體檢測及例項分割實驗

作者進一步在更復雜的目標檢測和例項分割上驗證所提出的運算元的有效性。他們的實驗模型是用 X-volution 增強的 ResNet-50,具體是替換了 ResNet-50 最後一個階段的三個瓶頸單元。為了充分的對比,作者展示了兩種形態的 X-volution,如表 2 和表 3 所示:X-volution(SA) 表示的是卷積與 global self-attention 結合,這種形態是為了驗證 X-volution 採用的卷積和 self-attention 結合模式的可行性;X-volution 則表示卷積和 PSSA 直接結合的形式,其為了檢驗所提出的 PSSA 的可行性。

表 2:MS COCO 物體檢測實驗結果

表 3:MS COCO 例項分割實驗結果

從表 2 與表 3 可以看出,兩種 X-volution 模式都獲得了大幅度的效能提升。其中,X-volution(SA) 更為明顯,這說明卷積和 self-attention 的互補對效能提升具有重大意義。而採用 PSSA 的 X-volution 效能也非常不俗,基本與 self-attention 效能相當,這也驗證了採用 PSSA 逼近 self-attention 是有效且合理的。

消融實驗

最後,作者詳細研究了卷積部分和 self-attention 部分對於運算元效能的影響。他們改變卷積部分的濾波核的尺寸大小,其效能變化結果如圖 2(a)。可以看到當卷積核設定為 1 時候,單純的卷積運算元效能退化非常嚴重,而此時的 X-volution 依然能取得 39.1 的平均準確率。當卷積核逐步增大後,純卷積的網路效能先升後降,而對應的 X-volution 運算元也呈同樣趨勢,但一直保持著對於單純卷積的效能優勢。其中,當卷積核為 7 時候效能最好。從這個結果可以得知,卷積在兩者間起到了較大的作用,這一點也與 AA-Net 結論相同。由於學習參數較多,卷積的作用也比較突出。

圖 2,(a) 研究不同卷積核對於效能的影響;(b) 不同形態運算元的優化效能比較。

在圖 2(b) 中,研究者展示了卷積、PSSA 和 X-volution 三種運算元在 MS COCO 目標檢測上的收斂曲線。可以看到,卷積在最開始時效能優於 X-volution 和 PSSA;而經過 3 個週期訓練後,X-volution 開始超越卷積,但是作為self-attention的逼近形式,PSSA在前9個週期收斂效能稍弱於卷積。其後,X-volution 一直保持顯著的領先。PSSA 則在 10 個週期後與卷積相當或略好於卷積。這組曲線證實了卷積的局部歸納偏置對於其訓練是有明顯的幫助,而低偏置的 self-attention 則收斂顯著慢於卷積,但經過較長時間訓練後可以超越卷積。將兩者整合的 X-volution 則兼備了卷積和 self-attention 的特性,同時展現了優秀的優化特性和良好的效能。

總結

作者在文中提出了一種新型的運算元——X-volution,整合了卷積和 self-attention 的互補特性。同時,他們從 self-attention 的公式中匯出了一種巧妙的全局自注意力的逼近形式——PSSA。作者通過在分類、目標檢測、例項分割等任務中的優秀表現證實了所提出的運算元的有效性。實驗也揭露了卷積與 self-attention 的配合確實能較為顯著地提升效能,並且兩者達到了實質上的特性互補。該新型運算元在效能上的顯著提升、計算的通用性與即插即用性方面的巨大優勢,為深度學習基礎計算單元庫,以及 NPU 計算架構的演進提供了一種新的基礎武器。

注:該論文作者是上海交通大學海思實習生陳炫宏和王航,由計算機視覺知名專家倪冰冰教授指導。


IT145.com E-mail:sddin#qq.com