MSViT: Dynamische Mixed-Scale-Tokenisierung für Vision Transformers

MSViT: Dynamic Mixed-Scale Tokenization for Vision Transformers

July 5, 2023

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

cs.AI

Zusammenfassung

Die Eingabe-Tokens für Vision Transformers tragen wenig semantische Bedeutung, da sie als regelmäßige, gleich große Bildausschnitte des Eingabebildes definiert sind, unabhängig von dessen Inhalt. Die Verarbeitung gleichmäßiger Hintergrundbereiche eines Bildes sollte jedoch nicht so viel Rechenleistung erfordern wie dichte, unübersichtliche Bereiche. Um dieses Problem zu lösen, schlagen wir ein dynamisches Mixed-Scale-Tokenisierungsschema für ViT, MSViT, vor. Unsere Methode führt einen bedingten Gating-Mechanismus ein, der die optimale Token-Skala für jede Bildregion auswählt, sodass die Anzahl der Tokens dynamisch pro Eingabe bestimmt wird. Das vorgeschlagene Gating-Modul ist ressourcenschonend, unabhängig von der Wahl des Transformer-Backbones und innerhalb weniger Epochen (z. B. 20 Epochen auf ImageNet) mit geringem Trainingsaufwand trainierbar. Zusätzlich führen wir eine neuartige Verallgemeinerung des Batch-Shaping-Loss ein, um das bedingte Verhalten des Gates während des Trainings zu verbessern. Wir zeigen, dass unser Gating-Modul in der Lage ist, trotz lokaler Operation auf grober Patch-Ebene, sinnvolle Semantik zu erlernen. Wir validieren MSViT in den Aufgaben der Klassifikation und Segmentierung, wo es zu einem verbesserten Kompromiss zwischen Genauigkeit und Komplexität führt.

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.