通过推理时间注意力工程改善图像聚类并减少伪影

本文的目标是在无需重新训练或微调的情况下，提高预训练视觉变换器（ViT）模型，特别是DINOv2，在图像聚类任务中的性能。随着模型规模的增大，多头注意力机制的高范数异常现象在某些补丁中出现。我们观察到这种异常导致了零样本图像聚类准确率的下降。这些异常特征表现为注意力图中某些补丁标记的值远大于其他补丁标记的值。为了解决这些问题，我们提出了一种称为推理时注意力工程（Inference-Time Attention Engineering, ITAE）的方法，该方法在推理过程中对注意力函数进行操作。具体而言，我们通过研究多头注意力机制中的一个查询-键-值（Query-Key-Value, QKV）补丁来识别这些异常，并减弱其在预训练模型中的相应注意力值。ITAE在多个数据集上展示了更高的聚类准确率，因为它在潜在空间中生成了更具表现力的特征。我们的研究结果突显了ITAE作为一种实用解决方案的潜力，可以在不重新训练或微调的情况下减少预训练ViT模型中的异常现象并提升其在聚类任务中的性能。