基于DINO知識蒸餾架構的分層級聯(lián)Transformer網(wǎng)絡

來自復旦大學、上海市智能信息處理重點實驗室和香港大學的研究者提出了一種基于 DINO 知識蒸餾架構的分層級聯(lián) Transformer (HCTransformer) 網(wǎng)絡。

小樣本學習是指從非常少量的標記數(shù)據(jù)中進行學習的問題，它有望降低標記成本，實現(xiàn)低成本、快速的模型部署，縮小人類智能與機器模型之間的差距。小樣本學習的關鍵問題是如何高效地利用隱藏在標注數(shù)據(jù)中的豐富信息中進行學習。近年來，視覺 Transformer (ViT [1]) 作為計算機視覺領域的新興架構，因其在很多主流任務中都取得了反超 CNN 的表現(xiàn)而受到廣泛關注。我們注意到 ViT 在訓練時易于陷入過擬合的狀態(tài)而表現(xiàn)出較差的性能，現(xiàn)有的研究大部分關注其在大、中型數(shù)據(jù)集上的表現(xiàn)，而幾乎沒有將其應用于針對更加敏感于過擬合的小樣本任務上的相關研究。 ?

為了緩解這一問題，提高小數(shù)據(jù)集下數(shù)據(jù)的利用效率，來自復旦大學、上海市智能信息處理重點實驗室和香港大學的研究者提出了一種基于 DINO [2] 知識蒸餾架構的分層級聯(lián) Transformer (HCTransformer) 網(wǎng)絡，通過譜聚類 tokens 池化以利用圖像的內(nèi)在結構來減少前景內(nèi)容和背景噪聲的模糊性，同時也提供了對不同數(shù)據(jù)集對于 patch token 切分大小的不同需求的適用性選擇；并利用一種非傳統(tǒng)的監(jiān)督方式，通過標簽的潛在屬性在圖像標簽中找到更豐富的視覺信息，而非簡單地學習由標簽分類的視覺概念。實驗表明，本文的方法在各個流行的 Few-Shot Classification 數(shù)據(jù)集上均取得了 SOTA 性能。 ?

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

代碼鏈接：https://github.com/StomachCold/HCTransformers

目前，這項研究已被 CVPR2022 接收，完整訓練代碼及模型已經(jīng)開源。 ? 方法

圖 1：方法整體架構。包括三個級聯(lián)的 Transformer 網(wǎng)絡，每一階段都是一個 teacher-student 的知識蒸餾結構。 ? DINO 自監(jiān)督架構 ? DINO 是一種基于知識蒸餾架構的自監(jiān)督網(wǎng)絡，和 BYOL [3] 類似，但 DINO 使用 Transformer 代替了 BYOL 中的 CNN。其中 Student 模型和 Teacher 模型的結構是完全一樣的。 ? 在訓練過程中，Teacher 模型并不通過輸入來更新參數(shù)，它的權重更新是通過對 Student 模型權重的 EMA（exponential moving average ） 來更新的。此外，Teacher 和 Student 都有獨立的數(shù)據(jù)增強方式。在這種設定下，整個網(wǎng)絡會把 Teacher 網(wǎng)絡的輸出當作是 Student 網(wǎng)絡的標簽，從而指導全局更新。因為網(wǎng)絡沒有直接使用圖片的類別標簽，所以可以支持使用較高維度的特征（如 65536）而不必局限于圖像的類別數(shù)量，使得網(wǎng)絡可以學習到更加細膩的特征表示。但在自監(jiān)督的架構下，由于小樣本學習的數(shù)據(jù)集較小，直接應用 DINO 進行訓練很難取得較好的效果。 ?

屬性代理監(jiān)督

基于這一前提，本文設法將 DINO 改成有監(jiān)督模型，但 Transformer 在小數(shù)據(jù)集上直接使用 one-hot 向量作為標簽，不僅需要將網(wǎng)絡輸出維度改成圖片類別數(shù)導致學習相對粗糙，而且會產(chǎn)生較嚴重的過擬合現(xiàn)象。因此，本文針對小樣本任務，為 Transformer 設計了一種屬性代理監(jiān)督的傳播方案，以避免傳統(tǒng)監(jiān)督帶來的局限性。 ? 具體來說，對于標簽空間中的每個視覺概念 y，我們的目標是為其學習一個語義屬性的代理向量，進而使同一個標簽的輸出均向這個代理向量在向量空間中靠攏，也同時影響到該代理向量的更新。其中，y 是輸入數(shù)據(jù)的標簽，D' 是整個網(wǎng)絡的輸出維度。 ?

在 Transformer 中，輸入圖片在通過 Encoder 后會得到一個 class token 和 N 個 patch tokens。一般來說在分類任務中，只會單一使用某一種token作為最終輸出結果進行監(jiān)督，而丟棄另一種token。相較于傳統(tǒng)設計，本文將 class token 和 patch tokens 同時都用于監(jiān)督，以更大限度提高數(shù)據(jù)利用率。

Class Token 監(jiān)督

圖片輸入到 DINO 網(wǎng)絡中，會經(jīng)過 encoder 后生成一個 class token 和 N 個的 patch tokens，之后 class token 再通過一個 MLP 得到最終的輸出。不同于傳統(tǒng)監(jiān)督方式，由于 DINO 本身的自監(jiān)督設計，導致最終的輸出并不是一個 one-hot 向量，而是一個更高維度的向量。因此我們?yōu)槊恳粋€標簽設計了一個維度的可學習向量作為該類別的代理屬性，通過 KL 散度將同一標簽的輸出聚攏。

在公式 1 中，是 student 網(wǎng)絡中的 class token 經(jīng)過 MLP 的最終輸出，是該標簽對應的語義屬性代理向量。為 KL 散度計算。 ? Patch Tokens 監(jiān)督 ? 在 Transformer 中，由于缺少 patch 級別的標簽信息，patch tokens 一般很難與 class token 同時用來監(jiān)督網(wǎng)絡。為了對 patch tokens 進行監(jiān)督，我們把 N 個 patch tokens 用經(jīng)過 softmax 操作后的注意力矩陣進行加權合并計算，得到一個全局 token，然后采用和 class token 相同的監(jiān)督方式對該全局 token 進行監(jiān)督。 ?

其中，為 class token 對其他 patch tokens 的注意力矩陣，是 encoder 后輸出的 patch tokens，為加權合并后的全局 token。 ?

公式 3 與公式 1 采用相同的策略。 ? 譜聚類 tokens 池化 ? 許多工作, 如 GroupFPN [4] 和 GLOM [5] 都證明了多尺度層次結構對 CNN 和 Transformer 都有效。本文的設計目的是為了將層次結構嵌入到 Transformer 中來提高網(wǎng)絡對特征的判別能力。與 Swin Transformer [4] 使用固定的網(wǎng)格池化方案不同，本文利用不規(guī)則網(wǎng)格池化方法來更靈活地匹配圖像結構。由于 Transformer 將在 tokens 之間生成自注意力矩陣，因此它為譜聚類算法 [5] 提供了強大的先驗經(jīng)驗，以根據(jù)語義相似性和空間布局對 tokens 進行分割合并。因此我們提出了一種基于譜聚類的池化方法，稱為譜聚類 tokens 池化。 ? 為了保持各 patch 之間原有的位置關系信息，本文對于 ViT 中的 N 個 patch tokens 之間的注意力矩陣加上了鄰接限制，使每個 patch 只和相鄰的 8 個 patch 有相似性。 ?

其中，A_p 是 patch tokens 的注意力矩陣，H 是包含位置信息的鄰接矩陣。 然后我們對 S 矩陣做一個 softmax 操作得到我們最終需要的 S' 矩陣作為譜聚類 tokens 池化的輸入。 ?

如上算法所示，每一次池化后我們將得到同類聚類數(shù)量減半的新 tokens。 ? 在本文的網(wǎng)絡設計中，訓練階段首先會訓練第一階段的 Transformer 網(wǎng)絡來得到一個有較好表征能力的特征生成器，隨后再引入池化操作，加上二三階段 Transformer 共同訓練。最終在三階段中通過驗證集挑選出最優(yōu)結果作為最終輸出。 ?

圖 2：譜聚類池化的可視化圖 ?

實驗結果

我們分別在四個流行的 Few-Shot Classification 數(shù)據(jù)集：mini-Imagenet, tiered-Imagenet, CIFAR-FS 和 FC100 上做了詳盡的實驗。 ?

最終結果如表 1,2,3 所示：相比于現(xiàn)有的 SOTA 模型，HCTransformers 在 1-shot 和 5-shot 的結果上都顯示出明顯的性能優(yōu)勢。例如，如表 1 所示，在 miniImagnet 上，HCTransformers 比最優(yōu) SOTA 結果分別高出 5.37%(1-shot) 和 4.03%(5-shot)。在 tieredImagenet 上，我們的方法在 1-shot 和 5-shot 上的表現(xiàn)分別比最好的 DC [6] 方法高出 1.48% 和 1.81%。與 DC 相比，我們不需要從 base 訓練集中借用類某類圖像的統(tǒng)計信息，并且使用更輕量級的分類器。此外，我們的方法和第三好的方法之間的差距是 5.09%，這也進一步驗證了我們的貢獻。 ? 這樣令人印象深刻的結果要歸功于我們的網(wǎng)絡結構，它能學習到數(shù)據(jù)中固有的內(nèi)在信息，并具有良好的泛化能力。表 2 和表 3 分別顯示了在小分辨率數(shù)據(jù)集 CIFAR-FS 和 FC100 上的結果。HCTransformers 在這些低分辨率設置中顯示出和 SOTA 類似或更好的結果: 在 CIFAR-FS 上提高了 1.02%(1-shot) 和 0.76%(5-shot); 在 FC100 上提高了 0.51%(1-shot) 和 1.12%(5-shot)。在小分辨率數(shù)據(jù)集上，我們并沒有超越以前的 SOTA 方法很多，我們將其歸因于 ViT 的 patch 機制：當圖像分辨率較小時，如 32*32，每一個 patch 所包含的實際像素過少，很難提取出有用的特征表示。DeepEMD [7] 的實驗也佐證了 patch cropping 會對小分辨率圖像產(chǎn)生負面影響。即便如此，我們的方法仍然在這兩個基準上都取得了新的 SOTA 結果。 ?
參考文獻
[1]. Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov,Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner,Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An imageis worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021.OpenReview.net,2021. 1, 2, 3, 7, 8
[2]. Mathilde Caron, Hugo Touvron, Ishan Misra, Herv′e J′egou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. arXiv preprint arXiv:2104.14294, 2021. 3, 5
[3]. Jean-Bastien Grill, Florian Strub, Florent Altch′e, Corentin Tallec, Pierre H.Richemond, Elena Buchatskaya, Carl Doersch, Bernardo ′Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukcuoglu, R′emi Munos, and Michal Valko. Bootstrap your own latent-A new approach to self-supervised learning. In Hugo Larochelle, Marc’Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. 3
[4]. Gangming Zhao, Weifeng Ge, and Yizhou Yu. Graphfpn:Graph feature pyramid network for object detection. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 2763–2772, 2021. 1
[5]. Geoffrey Hinton. How to represent part-whole hierarchies in a neural network. arXiv preprint arXiv:2102.12627, 2021. 1
[6]. Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. International Conference on Computer Vision (ICCV), 2021. 2, 4
[7]. Andrew Y Ng, Michael I Jordan, and Yair Weiss. On spectral clustering: Analysis and an algorithm. In Advances in neural information processing systems, pages 849–856, 2002. 1, 5
[8]. Shuo Yang, Lu Liu, and Min Xu. Free lunch for few-shot learning: Distribution calibration. In International Conference on Learning Representations (ICLR), 2021. 6
[9]. Chi Zhang, Yujun Cai, Guosheng Lin, and Chunhua Shen. Deepemd: Few-shot image classification with differentiable earth mover’s distance and structured classifiers. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020. 6, 7, 8