DeiT-LT Distillation Strikes Back for Vision Transformer Training on Long-Tailed Datasets

2024年04月03日
Vision Transformer (ViT)已成为各种计算机视觉任务中突出的架构。在ViT中,我们将输入图像分成补丁标记,并通过一堆自我关注块进行处理。然而,与卷积神经网络(CNN)不同,ViT的简单架构没有信息归纳偏差(例如局部性等)。因此,ViT需要大量的数据进行预训练。已经提出了各种数据有效的方法(DeiT)来有效地在平衡的数据集上训练ViT。然而,有限的文献讨论了ViT在长尾不平衡数据集上的使用。在这项工作中,我们介绍了DeiT-LT来解决从头开始训练ViTs的长尾数据集的问题。在DeiT-LT中,我们通过使用超出分布图像并重新加权蒸馏损失来引入了一种有效的从CNN蒸馏DIST标记的方式,从而解决了该问题。这导致早期ViT块中学习类似本地CNN的特征,提高了尾部类的泛化能力。此外,为了减少过度拟合,我们建议从平坦的CNN教师蒸馏,这导致学习DIST标记在所有ViT块中的低秩可推广特征。通过所提出的DeiT-LT方案,蒸馏DIST标记成为尾部类的专家,分类器CLS标记成为头部类的专家。专家有助于使用同一ViT架构中的不同标记集有效地学习对应于多数和少数类的特征。我们展示了DeiT-LT在从小规模CIFAR-10 LT到大规模iNaturalist-2018的数据集上从头开始训练ViT的有效性。
Vision Transformer (ViT) has emerged as a prominent architecture for various computer vision tasks. In ViT, we divide the input image into patch tokens and process them through a stack of self attention blocks. However, unlike Convolutional Neural Networks (CNN), ViTs simple architecture has no informative inductive bias (e.g., locality,etc. ). Due to this, ViT requires a large amount of data for pre-training. Various data efficient approaches (DeiT) have been proposed to train ViT on balanced datasets effectively. However, limited literature discusses the use of ViT for datasets with long-tailed imbalances. In this work, we introduce DeiT-LT to tackle the problem of training ViTs from scratch on long-tailed datasets. In DeiT-LT, we introduce an efficient and effective way of distillation from CNN via distillation DIST token by using out-of-distribution images and re-weighting the distillation loss to enhance focus on tail classes. This leads to the learning of local CNN-like features in early ViT blocks, improving generalization for tail classes. Further, to mitigate overfitting, we propose distilling from a flat CNN teacher, which leads to learning low-rank generalizable features for DIST tokens across all ViT blocks. With the proposed DeiT-LT scheme, the distillation DIST token becomes an expert on the tail classes, and the classifier CLS token becomes an expert on the head classes. The experts help to effectively learn features corresponding to both the majority and minority classes using a distinct set of tokens within the same ViT architecture. We show the effectiveness of DeiT-LT for training ViT from scratch on datasets ranging from small-scale CIFAR-10 LT to large-scale iNaturalist-2018.
许愿