MSViT : Tokenisation Dynamique à Échelles Mixtes pour les Transformers Visuels

MSViT: Dynamic Mixed-Scale Tokenization for Vision Transformers

July 5, 2023

Auteurs: Jakob Drachmann Havtorn, Amelie Royer, Tijmen Blankevoort, Babak Ehteshami Bejnordi

cs.AI

Résumé

Les tokens d'entrée des Vision Transformers portent peu de sens sémantique, car ils sont définis comme des patches réguliers de taille égale de l'image d'entrée, indépendamment de son contenu. Cependant, le traitement des zones d'arrière-plan uniformes d'une image ne devrait pas nécessiter autant de calcul que les zones denses et encombrées. Pour résoudre ce problème, nous proposons un schéma de tokenisation dynamique à échelles mixtes pour les ViT, appelé MSViT. Notre méthode introduit un mécanisme de gating conditionnel qui sélectionne l'échelle optimale des tokens pour chaque région de l'image, de sorte que le nombre de tokens soit déterminé dynamiquement par entrée. Le module de gating proposé est léger, indépendant du choix de l'architecture de transformer sous-jacente, et s'entraîne en quelques époques (par exemple, 20 époques sur ImageNet) avec un faible surcoût d'entraînement. De plus, pour améliorer le comportement conditionnel du gating pendant l'entraînement, nous introduisons une nouvelle généralisation de la fonction de perte de mise en forme par lot. Nous montrons que notre module de gating est capable d'apprendre des sémantiques significatives malgré son fonctionnement local au niveau des patches grossiers. Nous validons MSViT sur les tâches de classification et de segmentation, où il permet d'améliorer le compromis précision-complexité.

English

The input tokens to Vision Transformers carry little semantic meaning as they are defined as regular equal-sized patches of the input image, regardless of its content. However, processing uniform background areas of an image should not necessitate as much compute as dense, cluttered areas. To address this issue, we propose a dynamic mixed-scale tokenization scheme for ViT, MSViT. Our method introduces a conditional gating mechanism that selects the optimal token scale for every image region, such that the number of tokens is dynamically determined per input. The proposed gating module is lightweight, agnostic to the choice of transformer backbone, and trained within a few epochs (e.g., 20 epochs on ImageNet) with little training overhead. In addition, to enhance the conditional behavior of the gate during training, we introduce a novel generalization of the batch-shaping loss. We show that our gating module is able to learn meaningful semantics despite operating locally at the coarse patch-level. We validate MSViT on the tasks of classification and segmentation where it leads to improved accuracy-complexity trade-off.