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

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

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

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

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

文章中對以下三個核心問題進(jìn)行了研究：

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

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

第三，是否需要像 CNN 一樣進(jìn)行結(jié)構(gòu)化 Drop 操作？該方法嘗試了基于塊窗口和交叉窗口的結(jié)構(gòu)化 Drop 方式，并發(fā)現(xiàn)這種技巧對于 Vision Transformer 來說并不重要。

背景

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

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

DropKey

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

圖 2 DropKey 實(shí)現(xiàn)方法

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

雖然在 CNN 中對結(jié)構(gòu)化 drop 方法已有較為詳細(xì)的研究，但還沒有研究該 drop 方式對 ViT 的性能影響。為探究該策略會不會進(jìn)一步提升性能，該論文實(shí)現(xiàn)了 DropKey 的兩種結(jié)構(gòu)化形式，即 DropKey-Block 和 DropKey-Cross。其中，DropKey- Block 通過對以種子點(diǎn)為中心的正方形窗口內(nèi)連續(xù)區(qū)域進(jìn)行 drop，DropKey-Cross 則通過對以種子點(diǎn)為中心的十字形連續(xù)區(qū)域進(jìn)行 drop，如圖 3 所示。然而，該研究發(fā)現(xiàn)結(jié)構(gòu)化 drop 方法并不會帶來性能提升。

圖 3 DropKey 的結(jié)構(gòu)化實(shí)現(xiàn)方法

實(shí)驗(yàn)結(jié)果

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

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

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

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

圖 8 不同結(jié)構(gòu)化 drop 方法的性能比較

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

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

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

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

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

總結(jié)

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