ReMixer

Abstract

Vision Transformers (ViTs) have shown impressive results on various visual recognition tasks, alternating classic convolutional networks. While the initial ViTs treated all patches equally, recent studies reveal that incorporating inductive biases such as spatiality benefits the learned representations. However, most prior works solely focused on the location of patches, overlooking the scene structure of images. This paper aims to further guide the interaction of patches using the object information. Specifically, we propose ReMixer, which reweights the patch mixing layers of ViT based on the patch-wise object labels obtained in unsupervised or weakly-supervised manners, i.e., no additional human-annotating cost is necessary. Using the object labels, we compute a reweighting mask with a learnable scale parameter that calibrates the patch interactions, e.g., attention map of self-attention. We demonstrate that ReMixer improves ViTs over various downstream tasks, including classification, multi-object recognition, and background robustness. Finally, we show that our idea also works for MLP-Mixer and ConvMixer, implying its generic applicability to patch-based models.