近期，基于 Transformer 的算法被廣泛應用于計算機視覺的各類任務中，但該類算法在訓練數據量較小時容易產生過擬合問題。現有 Vision Transformer 通常直接引入 CNN 中常用的 Dropout 算法作為正則化器，其在注意力權重圖上進行隨機 Drop 并為不同深度的注意力層設置統一的 drop 概率。盡管 Dropout 十分簡單，但這種 drop 方式主要面臨三個主要問題。

首先，在 softmax 歸一化后進行隨機 Drop 會打破注意力權重的概率分布并且無法對權重峰值進行懲罰，從而導致模型仍會過擬合于局部特定信息（如圖 1）。其次，網絡深層中較大的 Drop 概率會導致高層語義信息缺失，而淺層中較小的 drop 概率會導致過擬合于底層細節特征，因此恒定的 drop 概率會導致訓練過程的不穩定。最后，CNN 中常用的結構化 drop 方式在 Vision Transformer 上的有效性并不明朗。

圖 1 不同正則化器對注意力分布圖的影響

美圖影像研究院（MT Lab）與中國科學院大學在 CVPR 2023 上發表了一篇文章，提出一種新穎且即插即用的正則化器 DropKey，該正則化器可以有效緩解 Vision Transformer 中的過擬合問題。

論文鏈接：https://arxiv.org/abs/2208.02646

文章中對以下三個核心問題進行了研究：

第一，在注意力層應該對什么信息執行 Drop 操作？與直接 Drop 注意力權重不同，該方法在計算注意力矩陣之前執行 Drop 操作，并將 Key 作為基礎 Drop 單元。該方法在理論上驗證了正則化器 DropKey 可以對高注意力區域進行懲罰并將注意力權值分配到其它感興趣的區域，從而增強模型對全局信息的捕獲能力。

第二，如何設置 Drop 概率？與所有層共享同一個 Drop 概率相比，該論文提出了一種新穎的 Drop 概率設置方法，即隨著自注意力層的加深而逐漸衰減 Drop 概率值。

第三，是否需要像 CNN 一樣進行結構化 Drop 操作？該方法嘗試了基于塊窗口和交叉窗口的結構化 Drop 方式，并發現這種技巧對于 Vision Transformer 來說并不重要。

背景

Vision Transformer（ViT）是近期計算機視覺模型中的新范式，它被廣泛地應用于圖像識別、圖像分割、人體關鍵點檢測和人物互相檢測等任務中。具體而言，ViT 將圖片分割為固定數量的圖像塊，將每個圖像塊都視作一個基本單位，同時引入了多頭自注意力機制來提取包含相互關系的特征信息。但現有 ViT 類方法在小數據集上往往會出現過擬合問題，即僅使用目標局部特征來完成指定任務。

為了克服以上問題，該論文提出了一種即插即拔、僅需要兩行代碼便可實現的正則化器 DropKey 用以緩解 ViT 類方法的過擬合問題。不同于已有的 Dropout，DropKey 將 Key 設置為 drop 對象并從理論和實驗上驗證了該改變可以對高注意力值部分進行懲罰，同時鼓勵模型更多關注與目標有關的其他圖像塊，有助于捕捉全局魯棒特征。此外，該論文還提出為不斷加深的注意力層設置遞減的 drop 概率，這可以避免模型過度擬合低級特征并同時保證有充足的高級特征以進行穩定的訓練。此外，該論文還通過實驗證明，結構化 drop 方法對 ViT 來說不是必要的。

DropKey

為了探究引發過擬合問題的本質原因，該研究首先將注意力機制形式化為一個簡單的優化目標并對其拉格朗日展開形式進行分析。發現當模型在不斷地優化時，當前迭代中注意力占比越大的圖像塊，在下次迭代過程中會傾向于被分配更大的注意力權值。為緩解這一問題，DropKey 通過隨機 drop 部分 Key 的方式來隱式地為每個注意力塊分配一個自適應算子以約束注意力分布從而使其變得更加平滑。值得注意的是，相對于其他根據特定任務而設計的正則化器，DropKey 無需任何手工設計。由于在訓練階段對 Key 執行隨機 drop，這將導致訓練和測試階段的輸出期望不一致，因此該方法還提出使用蒙特卡洛方法或微調技巧以對齊輸出期望。此外，該方法的實現僅需兩行代碼，具體如圖 2 所示。

圖 2 DropKey 實現方法

一般而言，ViT 會疊加多個注意力層以逐步學習高維特征。通常，較淺層會提取低維視覺特征，而深層則旨在提取建模空間上粗糙但復雜的信息。因此，該研究嘗試為深層設置較小的 drop 概率以避免丟失目標對象的重要信息。具體而言，DropKey 并不在每一層以固定的概率執行隨機 drop，而是隨著層數的不斷加深而逐漸降低 drop 的概率。此外，該研究還發現這種方法不僅適用于 DropKey，還可以顯著提高 Dropout 的性能。

雖然在 CNN 中對結構化 drop 方法已有較為詳細的研究，但還沒有研究該 drop 方式對 ViT 的性能影響。為探究該策略會不會進一步提升性能，該論文實現了 DropKey 的兩種結構化形式，即 DropKey-Block 和 DropKey-Cross。其中，DropKey- Block 通過對以種子點為中心的正方形窗口內連續區域進行 drop，DropKey-Cross 則通過對以種子點為中心的十字形連續區域進行 drop，如圖 3 所示。然而，該研究發現結構化 drop 方法并不會帶來性能提升。

圖 3 DropKey 的結構化實現方法

實驗結果

圖 4 DropKey 和 Dropout 在 CIFAR10/100 上的性能比較

圖 5 DropKey 和 Dropout 在 CIFAR100 上的注意力圖可視化效果比較

圖 6 不同 drop 概率設置策略的性能比較

圖 7 不同輸出期望對齊策略的性能比較

圖 8 不同結構化 drop 方法的性能比較

圖 9 DropKey 和 Dropout 在 ImageNet 上的性能比較

圖 10 DropKey 和 Dropout 在 COCO 上的性能比較

圖 11 DropKey 和 Dropout 在 HICO-DET 上的性能比較

圖 12 DropKey 和 Dropout 在 HICO-DET 上的性能比較

圖 13 DropKey 和 Dropout 在 HICO-DET 上的注意力圖可視化比較

總結

該論文創新性地提出了一種用于 ViT 的正則化器，用于緩解 ViT 的過擬合問題。與已有的正則化器相比，該方法可以通過簡單地將 Key 置為 drop 對象，從而為注意力層提供平滑的注意力分布。另外，該論文還提出了一種新穎的 drop 概率設置策略，成功地在有效緩解過擬合的同時穩定訓練過程。最后，該論文還探索了結構化 drop 方式對模型的性能影響。