正則化方法DropKey: 兩行代碼高效緩解視覺Transformer過(guò)擬合

美圖影像研究院（MT Lab）與中國(guó)科學(xué)院大學(xué)突破性地提出正則化方法 DropKey，用于緩解 Vision Transformer 中的過(guò)擬合問(wèn)題。該方法通過(guò)在注意力計(jì)算階段隨機(jī) drop 部分 Key 以鼓勵(lì)網(wǎng)絡(luò)捕獲目標(biāo)對(duì)象的全局信息，從而避免了由過(guò)于聚焦局部信息所引發(fā)的模型偏置問(wèn)題，繼而提升了基于 Transformer 的視覺類算法的精度。該論文已被計(jì)算機(jī)視覺三大頂會(huì)之一 CVPR 2023 接收。

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

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

▲圖1. 不同正則化器對(duì)注意力分布圖的影響 美圖影像研究院（MT Lab）與中國(guó)科學(xué)院大學(xué)在 CVPR 2023 上發(fā)表了一篇文章，提出一種新穎且即插即用的正則化器 DropKey，該正則化器可以有效緩解 Vision Transformer 中的過(guò)擬合問(wèn)題。

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

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

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

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

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

背景

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

為了克服以上問(wèn)題，該論文提出了一種即插即拔、僅需要兩行代碼便可實(shí)現(xiàn)的正則化器 DropKey 用以緩解 ViT 類方法的過(guò)擬合問(wèn)題。不同于已有的 Dropout，DropKey 將 Key 設(shè)置為 drop 對(duì)象并從理論和實(shí)驗(yàn)上驗(yàn)證了該改變可以對(duì)高注意力值部分進(jìn)行懲罰，同時(shí)鼓勵(lì)模型更多關(guān)注與目標(biāo)有關(guān)的其他圖像塊，有助于捕捉全局魯棒特征。

此外，該論文還提出為不斷加深的注意力層設(shè)置遞減的 drop 概率，這可以避免模型過(guò)度擬合低級(jí)特征并同時(shí)保證有充足的高級(jí)特征以進(jìn)行穩(wěn)定的訓(xùn)練。此外，該論文還通過(guò)實(shí)驗(yàn)證明，結(jié)構(gòu)化 drop 方法對(duì) ViT 來(lái)說(shuō)不是必要的。

DropKey

為了探究引發(fā)過(guò)擬合問(wèn)題的本質(zhì)原因，該研究首先將注意力機(jī)制形式化為一個(gè)簡(jiǎn)單的優(yōu)化目標(biāo)并對(duì)其拉格朗日展開形式進(jìn)行分析。發(fā)現(xiàn)當(dāng)模型在不斷地優(yōu)化時(shí)，當(dāng)前迭代中注意力占比越大的圖像塊，在下次迭代過(guò)程中會(huì)傾向于被分配更大的注意力權(quán)值。為緩解這一問(wèn)題，DropKey 通過(guò)隨機(jī) drop 部分 Key 的方式來(lái)隱式地為每個(gè)注意力塊分配一個(gè)自適應(yīng)算子以約束注意力分布從而使其變得更加平滑。

值得注意的是，相對(duì)于其他根據(jù)特定任務(wù)而設(shè)計(jì)的正則化器，DropKey 無(wú)需任何手工設(shè)計(jì)。由于在訓(xùn)練階段對(duì) Key 執(zhí)行隨機(jī) drop，這將導(dǎo)致訓(xùn)練和測(cè)試階段的輸出期望不一致，因此該方法還提出使用蒙特卡洛方法或微調(diào)技巧以對(duì)齊輸出期望。此外，該方法的實(shí)現(xiàn)僅需兩行代碼，具體如圖 2 所示。

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

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

雖然在 CNN 中對(duì)結(jié)構(gòu)化 drop 方法已有較為詳細(xì)的研究，但還沒有研究該 drop 方式對(duì) ViT 的性能影響。為探究該策略會(huì)不會(huì)進(jìn)一步提升性能，該論文實(shí)現(xiàn)了 DropKey 的兩種結(jié)構(gòu)化形式，即 DropKey-Block 和 DropKey-Cross。

其中，DropKey- Block 通過(guò)對(duì)以種子點(diǎn)為中心的正方形窗口內(nèi)連續(xù)區(qū)域進(jìn)行 drop，DropKey-Cross 則通過(guò)對(duì)以種子點(diǎn)為中心的十字形連續(xù)區(qū)域進(jìn)行 drop，如圖 3 所示。然而，該研究發(fā)現(xiàn)結(jié)構(gòu)化 drop 方法并不會(huì)帶來(lái)性能提升。

▲圖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. 不同輸出期望對(duì)齊策略的性能比較

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

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

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

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

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

總結(jié)

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

審核編輯 ：李倩