[论文解读] CaFT: Clustering and Filter on Tokens of Transformer for Weakly Supervised Object Localization

Ming Li|arXiv (Cornell University)|Jan 3, 2022

Advanced Neural Network Applications被引用 8

一句话总结

CaFT 提出了一种新颖的弱监督目标定位方法，通过聚类 ViT 的 patch tokens 并应用轻量级注意力过滤（AtF）头，无需人工设定阈值即可生成精确的目标掩码。在使用真实类别标签的情况下，该方法在 CUB-200 和 ImageNet-1K 上分别实现了 97.55% 和 69.86% 的定位准确率，优于以往基于 CAM 的方法，其优势在于摆脱了对判别性区域的依赖以及对阈值选择的依赖。

ABSTRACT

Weakly supervised object localization (WSOL) is a challenging task to localize the object by only category labels. However, there is contradiction between classification and localization because accurate classification network tends to pay attention to discriminative region of objects rather than the entirety. We propose this discrimination is caused by handcraft threshold choosing in CAM-based methods. Therefore, we propose Clustering and Filter of Tokens (CaFT) with Vision Transformer (ViT) backbone to solve this problem in another way. CaFT first sends the patch tokens of the image split to ViT and cluster the output tokens to generate initial mask of the object. Secondly, CaFT considers the initial mask as pseudo labels to train a shallow convolution head (Attention Filter, AtF) following backbone to directly extract the mask from tokens. Then, CaFT splits the image into parts, outputs masks respectively and merges them into one refined mask. Finally, a new AtF is trained on the refined masks and used to predict the box of object. Experiments verify that CaFT outperforms previous work and achieves 97.55\% and 69.86\% localization accuracy with ground-truth class on CUB-200 and ImageNet-1K respectively. CaFT provides a fresh way to think about the WSOL task.

研究动机与目标

• 解决基于 CAM 的方法因人工设定阈值而仅关注判别性区域而非完整物体的局限性。
• 消除基于激活图的定位方法中对手动阈值选择的依赖。
• 通过利用 Vision Transformer 中的类别 token 实现基于聚类的掩码生成，提升定位准确率。
• 开发一种轻量级、伪监督的头部（AtF），以从聚类生成的初始掩码中优化掩码，实现鲁棒的目标定位。
• 通过将聚类与 ViT 特征上的深度学习相结合，为 WSOL 提供一种新范式。

提出的方法

• 对 ViT 处理后的 patch tokens 进行聚类，以分组特征向量，选择包含类别 token 的聚类作为前景物体掩码。
• 在初始聚类掩码作为伪标签的基础上训练一个轻量级注意力过滤（AtF）头部，以从 ViT tokens 中优化物体掩码。
• 将图像分割为多个部分，分别独立处理以生成局部掩码，随后将局部掩码合并为优化后的全局掩码。
• 在优化后的掩码上训练第二个 AtF 头，以预测最终的物体边界框。
• 在聚类和 AtF 之前，融合最后三层 ViT 的特征图及其位置嵌入，以提升表征质量。
• AtF 模块使用 1×1 或 3×3 卷积，配合批量归一化和激活函数，从 token 级特征中学习前景与背景的区分。

实验结果

研究问题

• RQ1 是否可以通过类别 token 引导的 ViT patch tokens 聚类，在无需阈值的情况下有效生成初始物体掩码？
• RQ2 在聚类生成的伪掩码上训练的轻量级注意力过滤（AtF）头部，是否能通过更好的稳定性和噪声过滤，优于直接回归或分类？
• RQ3 多尺度融合 ViT 特征（来自多个层）如何影响聚类和掩码生成的性能？
• RQ4 在避免过拟合并提升边缘检测能力的前提下，AtF 头的最佳架构（如卷积核大小与深度）是什么？
• RQ5 聚类生成的伪标签能否有效用于训练鲁棒的定位头部？与边界框回归相比表现如何？

主要发现

• CaFT 在使用真实类别标签的 CUB-200 上实现了 97.55% 的定位准确率，显著优于以往基于 CAM 的方法。
• 在 ImageNet-1K 上，CaFT 在使用真实类别标签时达到 69.86% 的定位准确率，展现出在大规模数据上的强大泛化能力。
• 在相同伪标签下，基于聚类生成伪掩码训练的 AtF 头优于边界框回归头，因其具备更好的稳定性和噪声过滤能力。
• 使用 1×1 卷积的 AtF（AtFS）性能最佳且泛化能力最强，而 3×3 卷积核引入过拟合并降低边缘检测性能。
• 融合最后三层 ViT 的特征可提升掩码质量，加入位置嵌入可进一步将性能提升 0.05%。
• 直接在伪掩码上对深层网络（如 ResNet50、VGG16）进行端到端训练会导致过拟合，性能反而劣于仅使用聚类的方法，凸显了轻量级、任务专用头部（如 AtF）的必要性。