Pas tous les patches sont ce dont vous avez besoin : Accélérer les Transformers visuels par le biais de réorganisations de tokens

Les Vision Transformers (ViTs) prennent en compte tous les patches d'image comme des tokens et construisent une attention multi-têtes auto (MHSA) entre eux. L'utilisation complète de ces tokens d'image entraîne des calculs redondants, car tous les tokens ne sont pas attentifs dans la MHSA. Par exemple, les tokens contenant des arrière-plans d'images sémantiquement sans sens ou distrayants n'apportent pas une contribution positive aux prédictions du ViT. Dans ce travail, nous proposons de réorganiser les tokens d'image au cours du processus de propagation avant des modèles ViT, ce qui est intégré à l'entraînement du ViT. Pour chaque inférence en avant, nous identifions les tokens d'image attentifs entre les modules MHSA et FFN (c'est-à-dire le réseau de propagation avant), guidés par l'attention du token de classe correspondant. Ensuite, nous réorganisons les tokens d'image en conservant ceux qui sont attentifs et en fusionnant ceux qui ne le sont pas pour accélérer les calculs ultérieurs de la MHSA et du FFN. Ainsi, notre méthode EViT améliore les ViTs sous deux angles. Premièrement, avec le même nombre de tokens d'image en entrée, notre méthode réduit les calculs de la MHSA et du FFN pour une inférence plus efficace. Par exemple, la vitesse d'inférence du DeiT-S est augmentée de 50 % tandis que sa précision de reconnaissance diminue seulement de 0,3 % pour la classification ImageNet. Deuxièmement, en maintenant le même coût computationnel, notre méthode permet aux ViTs d'utiliser un plus grand nombre de tokens d'image en entrée pour améliorer la précision de reconnaissance, où ces tokens proviennent d'images à plus haute résolution. Un exemple est que nous améliorons la précision de reconnaissance du DeiT-S de 1 % pour la classification ImageNet tout en maintenant le même coût computationnel qu'un DeiT-S standard. Parallèlement, notre méthode n'introduit pas davantage de paramètres aux ViTs. Les expériences sur des benchmarks standards montrent l'efficacité de notre méthode. Le code est disponible à l'adresse suivante : https://github.com/youweiliang/evit