MSViT: Tokenizzazione Dinamica a Scala Mista per Vision Transformers

MSViT: Dynamic Mixed-Scale Tokenization for Vision Transformers

July 5, 2023

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

cs.AI

Abstract

I token di input per i Vision Transformer trasportano poco significato semantico, poiché sono definiti come patch regolari di dimensioni uguali dell'immagine di input, indipendentemente dal suo contenuto. Tuttavia, elaborare aree uniformi di sfondo di un'immagine non dovrebbe richiedere la stessa quantità di calcolo rispetto ad aree dense e disordinate. Per affrontare questo problema, proponiamo uno schema di tokenizzazione dinamica a scala mista per ViT, denominato MSViT. Il nostro metodo introduce un meccanismo di gating condizionale che seleziona la scala ottimale dei token per ogni regione dell'immagine, in modo che il numero di token sia determinato dinamicamente per ogni input. Il modulo di gating proposto è leggero, indipendente dalla scelta del backbone del transformer, e viene addestrato in poche epoche (ad esempio, 20 epoche su ImageNet) con un sovraccarico di addestramento minimo. Inoltre, per migliorare il comportamento condizionale del gate durante l'addestramento, introduciamo una nuova generalizzazione della funzione di loss di batch-shaping. Dimostriamo che il nostro modulo di gating è in grado di apprendere semantiche significative nonostante operi localmente a livello di patch grossolane. Validiamo MSViT sui task di classificazione e segmentazione, dove porta a un miglior compromesso tra accuratezza e complessità.

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.