首次統(tǒng)一卷積與自注意力，上海交大、華為海思提出X-volution，發(fā)力網(wǎng)絡(luò)核心基礎(chǔ)架構(gòu)創(chuàng)新

卷積和自注意力各有優(yōu)勢，但二者的有效結(jié)合一直是一大難題。為了取二者之長，上海交大 - 華為海思聯(lián)合團(tuán)隊(duì)提出了一種名為 X-volution 的新型算子。該算子在性能上的顯著提升、計(jì)算的通用性與即插即用的特性為深度學(xué)習(xí)基礎(chǔ)計(jì)算單元庫以及 NPU 計(jì)算架構(gòu)的演進(jìn)提供了一種新的基礎(chǔ)武器。

眾所周知，卷積操作（convolution）與自注意力操作（self-attention）是深度學(xué)習(xí)兩大核心的基礎(chǔ)網(wǎng)絡(luò)計(jì)算單元（或稱為模型算子）。卷積操作通過線性乘子，提取圖像局部特征；自注意力操作通過高階乘子運(yùn)算，提取圖像全域 / 局部的特征關(guān)聯(lián)特性。兩種算子成為深度學(xué)習(xí)兩大重要網(wǎng)絡(luò)架構(gòu)演化——CNN 與 Transformer 的計(jì)算基石。兩種算子在圖像特征提取與語義抽象方面的互補(bǔ)性不言而喻：線性 vs. 高階， 局部 vs. 全局。因此，能否設(shè)計(jì)一種包含這兩種操作的融合算子并使其發(fā)揮互補(bǔ)優(yōu)勢，一直是深度學(xué)習(xí)架構(gòu)研究者熱衷的研究方向之一。

然而，由于卷積運(yùn)算與自注意力運(yùn)算在計(jì)算模式上的異構(gòu)性，這項(xiàng)任務(wù)存在巨大的挑戰(zhàn)。目前學(xué)界中的一些工作也在努力統(tǒng)一兩者，他們主要從拓?fù)浣Y(jié)構(gòu)組合角度來粗粒度地結(jié)合兩種算子，例如，發(fā)表在 ICCV 2019 上的 AA-Net 采用了一種將卷積中部分通道替換為由 self-attention 來處理，然后將卷積和 self-attention 分別處理的特征連接來達(dá)到聯(lián)合兩種算子的目的，這種做法證明了卷積和 self-attention 結(jié)合后確實(shí)能在分類、檢測、分割等基礎(chǔ)任務(wù)上達(dá)到比較可觀的性能收益。

然而，粗粒度的組合（本質(zhì)上就是兩路計(jì)算并聯(lián)）會導(dǎo)致其組合后網(wǎng)絡(luò)形態(tài)可用性下降。具體來說，卷積和 self-attention 運(yùn)算模式存在較大差異，兩者同時存在會導(dǎo)致網(wǎng)絡(luò)結(jié)構(gòu)不規(guī)則，進(jìn)而影響網(wǎng)絡(luò)推理效率，并不為目前一些工業(yè)界通用的芯片計(jì)算架構(gòu)所友好支持。同時組合后的算子在算力上也存在巨大的挑戰(zhàn)。

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

針對這些挑戰(zhàn)，日前，上海交大 - 華為海思聯(lián)合團(tuán)隊(duì)在 arXiv 上發(fā)表了「X-volution: On the Unification of Convolution and Self-attention」，首次在計(jì)算模式上統(tǒng)一了這兩大基礎(chǔ)算子，并在推理階段歸并成一個簡單的卷積型算子：X-volution。

X-volution 兼顧卷積與自注意力操作的互補(bǔ)優(yōu)勢，并且在現(xiàn)有通用網(wǎng)絡(luò)計(jì)算框架上不需要額外算子支持，也不增加除卷積外的額外算力或影響網(wǎng)絡(luò)的規(guī)范性 / 可用性（即插即用）。

該工作的突破主要受以下思路的啟發(fā)：對全局的 self-attention 進(jìn)行理論分析后，研究者發(fā)現(xiàn)在一定條件下（例如圖像 / 特征圖的鄰接像素滿足馬爾可夫性質(zhì)），全局的 self-attention 可以通過局部的 self-attention 配合卷積運(yùn)算來逼近。

具體來說，本文作者提出了一種新型的 self-attention 機(jī)制——PSSA。這種機(jī)制分為兩個步驟：首先將輸入的特征沿指定的多個方向進(jìn)行循環(huán)移位（采用索引來實(shí)現(xiàn)）得到移位后的特征，然后將移位后的特征與原特征通過元素點(diǎn)積獲得變換后的特征，再對該特征在局部區(qū)域進(jìn)行加權(quán)求和（可采用卷積來替代），至此獲得經(jīng)過注意力機(jī)制處理后的特征。通過層次堆疊，可以持續(xù)地將局部的上下文關(guān)系傳播到全局從而實(shí)現(xiàn)全局的 self-attention。

值得注意的是，PSSA 實(shí)際上將 self-attention 巧妙地轉(zhuǎn)化為了一個在簡單變換后的特征上的標(biāo)準(zhǔn)的卷積操作，這從形式上實(shí)現(xiàn)了 self-attention 向卷積的統(tǒng)一。利用此逼近式的 self-attention 機(jī)制，作者建立了一個多分枝的模塊將卷積和 self-attention 整合在一起，這個模塊從功能上實(shí)現(xiàn)了兩者的統(tǒng)一。

更重要的是，這個多分枝的結(jié)構(gòu)可以利用結(jié)構(gòu)重參數(shù)化的方法進(jìn)行有條件的合并。多分枝結(jié)構(gòu)可以合并為單個卷積，合并后可以獲得一個原子級的算子，稱為 X-volution（X-volution 的權(quán)重可以看作一個靜態(tài)卷積權(quán)重，以及一個內(nèi)容相關(guān)動態(tài)卷積權(quán)重的和）。此算子同時具備了卷積和 self-attention 的特性，且不會影響網(wǎng)絡(luò)的規(guī)范性 / 可用性。

作者在分類、檢測、分割等主流 SOTA 實(shí)驗(yàn)上取得了顯著的性能提升。

圖 1，算子詳細(xì)結(jié)構(gòu)框圖。受結(jié)構(gòu)重參數(shù)化思想啟發(fā)，X-volution 被設(shè)計(jì)為訓(xùn)練和推理階段結(jié)構(gòu)解耦的形式。它的訓(xùn)練結(jié)構(gòu)時有兩個主要分支（如中間所示），右分支由級聯(lián)的卷積和 BN 構(gòu)成，可為 X-volution 集成卷積的能力。左邊包括 PSSA，它提供近似的全局自注意力特性。完成訓(xùn)練后，X-volution 可以有條件地重新參數(shù)化為一個卷積操作。在推理階段，X-volution 實(shí)際上是一個動態(tài)卷積算子，它的權(quán)重包括需要在線計(jì)算的 attention 動態(tài)參數(shù)部分和已經(jīng)訓(xùn)練和固化的卷積靜態(tài)參數(shù)部分。

實(shí)驗(yàn)部分

作者將 X-volution 接入到經(jīng)典的 ResNet 模型中用于 ImageNet 分類、MS COCO 物體檢測、實(shí)例分割等關(guān)鍵基礎(chǔ)任務(wù)并都取得了不俗的提升。為了排除其他因素干擾，實(shí)驗(yàn)中作者所使用的 self-attention 和 PSSA 都沒有添加位置編碼，并且沒有對數(shù)據(jù)集進(jìn)行額外的增廣，沒有使用額外的訓(xùn)練技巧（如：余弦退火、標(biāo)簽平滑等）。

ImageNet 分類實(shí)驗(yàn)

對于 ImageNet 圖片分類實(shí)驗(yàn)，作者分別測試了在 ResNet 中三個不同位置接入 X-volution 的結(jié)果。將 X-volution 接入到常規(guī)的 ResNet 第五階段瓶頸單元的結(jié)果如表 1 所示：在 ResNet-34 與 ResNet-50 中均提升不明顯，這是因?yàn)樵诖穗A段的特征圖尺寸已經(jīng)接近卷積核大小。實(shí)驗(yàn)發(fā)現(xiàn)在第三階段效果最為突出，分別取得了 1.2% 與 0.9% 的顯著提升。值得注意的是，作者對于 ResNet 改動較小，但是性能卻依然能有大幅度的提升，這證實(shí)了文中所提出的 X-volution 算子具有良好的性能。

表 1. ImageNet 實(shí)驗(yàn)結(jié)果及瓶頸單元詳細(xì)結(jié)構(gòu)

MS COCO 物體檢測及實(shí)例分割實(shí)驗(yàn)

作者進(jìn)一步在更復(fù)雜的目標(biāo)檢測和實(shí)例分割上驗(yàn)證所提出的算子的有效性。他們的實(shí)驗(yàn)?zāi)Ｐ褪怯?X-volution 增強(qiáng)的 ResNet-50，具體是替換了 ResNet-50 最后一個階段的三個瓶頸單元。為了充分的對比，作者展示了兩種形態(tài)的 X-volution，如表 2 和表 3 所示：X-volution(SA) 表示的是卷積與 global self-attention 結(jié)合，這種形態(tài)是為了驗(yàn)證 X-volution 采用的卷積和 self-attention 結(jié)合模式的可行性；X-volution 則表示卷積和 PSSA 直接結(jié)合的形式，其為了檢驗(yàn)所提出的 PSSA 的可行性。

表 2：MS COCO 物體檢測實(shí)驗(yàn)結(jié)果

表 3：MS COCO 實(shí)例分割實(shí)驗(yàn)結(jié)果

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

消融實(shí)驗(yàn)

最后，作者詳細(xì)研究了卷積部分和 self-attention 部分對于算子性能的影響。他們改變卷積部分的濾波核的尺寸大小，其性能變化結(jié)果如圖 2(a)?？梢钥吹疆?dāng)卷積核設(shè)置為 1 時候，單純的卷積算子性能退化非常嚴(yán)重，而此時的 X-volution 依然能取得 39.1 的平均準(zhǔn)確率。當(dāng)卷積核逐步增大后，純卷積的網(wǎng)絡(luò)性能先升后降，而對應(yīng)的 X-volution 算子也呈同樣趨勢，但一直保持著對于單純卷積的性能優(yōu)勢。其中，當(dāng)卷積核為 7 時候性能最好。從這個結(jié)果可以得知，卷積在兩者間起到了較大的作用，這一點(diǎn)也與 AA-Net 結(jié)論相同。由于學(xué)習(xí)參數(shù)較多，卷積的作用也比較突出。

圖 2，(a) 研究不同卷積核對于性能的影響；(b) 不同形態(tài)算子的優(yōu)化性能比較。

在圖 2(b) 中，研究者展示了卷積、PSSA 和 X-volution 三種算子在 MS COCO 目標(biāo)檢測上的收斂曲線?？梢钥吹?，卷積在最開始時性能優(yōu)于 X-volution 和 PSSA；而經(jīng)過 3 個周期訓(xùn)練后，X-volution 開始超越卷積，但是作為self-attention的逼近形式，PSSA在前9個周期收斂性能稍弱于卷積。其后，X-volution 一直保持顯著的領(lǐng)先。PSSA 則在 10 個周期后與卷積相當(dāng)或略好于卷積。這組曲線證實(shí)了卷積的局部歸納偏置對于其訓(xùn)練是有明顯的幫助，而低偏置的 self-attention 則收斂顯著慢于卷積，但經(jīng)過較長時間訓(xùn)練后可以超越卷積。將兩者整合的 X-volution 則兼?zhèn)淞司矸e和 self-attention 的特性，同時展現(xiàn)了優(yōu)秀的優(yōu)化特性和良好的性能。

總結(jié)

作者在文中提出了一種新型的算子——X-volution，整合了卷積和 self-attention 的互補(bǔ)特性。同時，他們從 self-attention 的公式中導(dǎo)出了一種巧妙的全局自注意力的逼近形式——PSSA。作者通過在分類、目標(biāo)檢測、實(shí)例分割等任務(wù)中的優(yōu)秀表現(xiàn)證實(shí)了所提出的算子的有效性。實(shí)驗(yàn)也揭露了卷積與 self-attention 的配合確實(shí)能較為顯著地提升性能，并且兩者達(dá)到了實(shí)質(zhì)上的特性互補(bǔ)。該新型算子在性能上的顯著提升、計(jì)算的通用性與即插即用性方面的巨大優(yōu)勢，為深度學(xué)習(xí)基礎(chǔ)計(jì)算單元庫，以及 NPU 計(jì)算架構(gòu)的演進(jìn)提供了一種新的基礎(chǔ)武器。

注：該論文作者是上海交通大學(xué)海思實(shí)習(xí)生陳炫宏和王航，由計(jì)算機(jī)視覺知名專家倪冰冰教授指導(dǎo)。