Multi-Token Enhancing for Vision Representation Learning

TL;DR

This paper introduces Multi-Token Enhancing (MTE), a method that extracts multiple auxiliary tokens during training to improve vision representation learning without increasing inference costs, by distilling their knowledge into a global token.

Contribution

MTE is a novel approach that enhances self-supervised vision models by using auxiliary tokens during training and distilling their knowledge, avoiding additional inference costs.

Findings

Consistently improves performance across downstream tasks.

Compatible with various self-supervised loss functions and architectures.

Reduces additional inference costs through knowledge distillation.

Abstract

Vision representation learning, especially self-supervised learning, is pivotal for various vision applications. Ensemble learning has also succeeded in enhancing the performance and robustness of the vision models. However, traditional ensemble strategies are impractical for representation learning, especially self-supervised representation learning that requires large-scale datasets and long schedules. This is because they require k times more training and inference computation costs for an ensemble of k models. Differently, we introduce Multi-Token Enhancing (MTE) that extracts multiple auxiliary tokens simultaneously from a single model to enhance representation learning, while incurring minimal additional training costs and no additional inference costs. These auxiliary tokens, including auxiliary CLS tokens and adaptively pooled tokens, capture complementary information due to their differences. Meanwhile, to address the increase in inference costs, we distill the knowledge acquired by the auxiliary tokens into a global token during pre-training. Consequently, we can discard the auxiliary tokens during inference without incurring additional costs. Our MTE is compatible with various self-supervised loss functions and architectures, consistently improving performances across different downstream tasks. Our source code will be made publicly available.