實踐教程 | Evo-ViT：高性能Transformer加速方法

論文地址：

https://link.zhihu.com/?target=https%3A//arxiv.org/abs/2108.01390

前言

Transformer基礎網(wǎng)絡的高效性設計問題隨著Transformer在計算機視覺領域的蓬勃發(fā)展逐漸受到國內(nèi)外眾多學術機構和企業(yè)的關注。本文提出了一種新式的Transformer加速算法：Evo-ViT: Slow-Fast Token Evolution for Dynamic Vision Transformer，該算法能夠在保證分類準確率損失較小的情況下，大幅提升Transformer的推理速度，如在ImageNet 1K數(shù)據(jù)集下，Evo-ViT可以提升DeiT-S 60%推理速度的同時僅僅損失0.4%的精度。

研究意義與背景

最近，Vision Transformer 及其變體在各種計算機視覺任務中顯示出巨大的潛力。通過自注意力機制捕獲短程和長程視覺依賴的能力是其成功的主要來源。但是長程感受野同樣帶來了巨大的計算開銷，特別是對于高分辨率視覺任務（例如，目標檢測、分割）。研究者們開始研究如何在盡量保持原有模型準確率的前提下，降低模型計算復雜度，從而使得視覺 Transformer成為一種更加通用、高效、低廉的解決框架。

目前，主流的Transformer高效設計方案包括兩種：一種是借助結構化的空間先驗，如PiT[2]，LeViT[6]等利用空間下采樣構造金字塔型模型，再例如PVT[1]，Swin Transformer[9]等利用圖像局部先驗構造稀疏化的自注意力模塊；另一種是進行非結構化的網(wǎng)絡裁剪，例如DynamicViT[3]、PS-ViT[4]，基于預訓練好的模型，分析該模型的冗余性，對模型進行空間token或者特征通道的裁剪。然而，非結構化的裁剪會破化模型內(nèi)部特征的空間結構，使得這兩種方法無法相輔相成。

如圖1第三個分支，相比直接將信息量低的token裁剪掉，此次工作中提出一種新的即插即用的token雙流更新策略，能在模型訓練的同時動態(tài)判斷非結構性的token冗余及低信息分布，從而高效更新低信息token、精細更新高信息token，實現(xiàn)模型的高效準確建模，并保留了完整的空間結構。因此，該工作所提出的Evo-ViT方法可以同時適用于直筒型和金字塔型Transformer結構。

技術創(chuàng)新

與現(xiàn)有的Transformer高效設計方案相比，Evo-ViT是即插即用的加速策略，既適用于直筒型結構，也適用于金字塔型結構的視覺Transformer，不破壞原有模型的結構化設計；同時，Evo-ViT是在模型訓練過程中動態(tài)發(fā)掘冗余與低效信息，無需預訓練模型，因此能同時提升模型的訓練和推斷效率。方法主要兩點創(chuàng)新：

1. 提出了結構保留的token選擇策略，通過分析全局class attention，來動態(tài)區(qū)分高信息token和低信息token，并保留低信息token來確保完整的信息流；

2. 提出了雙流token更新策略，對高信息token及低信息token的歸納進行精細更新，然后用歸納token對低信息token進行高效更新，從而在不改變網(wǎng)絡結構的情況下，大幅提升模型性能。

技術細節(jié)

圖2 介紹了Evo-ViT的具體框架設計，包括基于全局class attention的token選擇以及慢速-快速雙流token更新兩個模塊。其根據(jù)全局class attention的排序判斷高信息token和低信息token，將低信息token整合為一個歸納token，和高信息token一起輸入到原始多頭注意力（Multi-head Self-Attention, MSA）模塊以及前向傳播（Fast Fed-forward Network, FFN）模塊中進行精細更新。更新后的歸納token用來快速更新低信息token。全局class attention也在精細更新過程中進行同步更新變化。

實驗結果

為了驗證方法的有效性，Evo-ViT基于直筒型Transformer結構DeiT[5]、金字塔型結構LeViT[6]，在主流Benchmark ImageNet-1k上進行對比實驗。

圖3是Evo-ViT和現(xiàn)有token相關高效性設計方法的對比，包括PS-ViT[4]、DynamicViT[3]、SViTE[7]、IA-RED2[8]。實驗結果表明，其在確保準確率的同時，能夠有更高的吞吐量提升，性能優(yōu)化的表現(xiàn)更佳。

圖4為Evo-ViT在金字塔型結構上的驗證實驗。由于以往直接裁剪的方法均無法直接用于具有空間先驗的金字塔型結構，Evo-ViT只和目前SOTA的金字塔型Transformer進行了比較。

該工作還進一步可視化了其token選擇結果，如圖5所示。左邊部分為訓練好的完整模型在各層的token選擇結果，右邊部分為訓練過程中不同階段的token選擇結果?？梢园l(fā)現(xiàn)，訓練好的模型各層選擇基本趨于一致，這是因為模型傾向于用更多的資源更新高信息量的token，即讓高信息量token通過所有層的精細更新。同時，由于文章提出的保持結構的雙流更新策略，可以發(fā)現(xiàn)一些淺層被誤判的token在深層也可以被撿回來，如第四行的棒球圖片，在前層時棒球桿被誤判為低信息量token，但是在深層全部被撿了回來。進一步觀察右邊部分，可以發(fā)現(xiàn)隨著訓練的深入token選擇效果逐漸趨于最優(yōu)。

結論

本文提出了一種基于慢速-快速雙流更新思想的通用視覺Transformer加速方法，Evo-ViT。不同于以往的方法，本文通過給高信息量token和低信息量token分配不同的計算優(yōu)先級，使得加速模型的同時保留了內(nèi)部特征的空間結構，同時適用于直筒型和金字塔型Transformer。實驗表明Evo-ViT可以對模型進行有效的加速。

從可視化結果可以看出，本文所提出的方法可以使模型更關注于圖像的核心區(qū)域，這對于模型的可解釋性，以及需要利用高層語義的任務有潛在的幫助。如何將本文的方法用于更多下游任務，如檢測、分割，也是一個有趣的方向。

以上即Evo-ViT的基本介紹，更多細節(jié)可見論文。大家有什么想法意見歡迎評論留言~

參考文獻

[1] Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. arXiv:2102.12122

[2] Rethinking spatial dimensions of vision transformers. arXiv preprint arXiv:2103.16302

[3] DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification. arXiv:2106.02034.

[4] Patch Slimming for Efficient Vision Transformers. arXiv:2106.02852

[5] Training data-efficient image transformers & distillation through attention. arXiv:2012.12877.

[6] LeViT: a Vision Transformer in ConvNet’s Clothing for Faster Inference. arXiv:2104.01136

[7] Chasing Sparsity in Vision Transformers: An End-to-End Exploration. arXiv:2106.04533

[8] IA-RED2: Interpretability-Aware Redundancy Reduction for Vision Transformers. arXiv:2106.12620

[9] Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2103.14030.