MSViT: Tokenização Dinâmica de Escalas Mistas para Vision Transformers

MSViT: Dynamic Mixed-Scale Tokenization for Vision Transformers

July 5, 2023

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

cs.AI

Resumo

Os tokens de entrada para os Vision Transformers carregam pouco significado semântico, pois são definidos como patches regulares de tamanho igual da imagem de entrada, independentemente de seu conteúdo. No entanto, processar áreas uniformes de fundo de uma imagem não deve exigir tanto poder computacional quanto áreas densas e desordenadas. Para resolver esse problema, propomos um esquema de tokenização dinâmica de escala mista para ViT, o MSViT. Nosso método introduz um mecanismo de portão condicional que seleciona a escala ótima de token para cada região da imagem, de modo que o número de tokens seja determinado dinamicamente por entrada. O módulo de portão proposto é leve, agnóstico à escolha do backbone do transformer e treinado em poucas épocas (por exemplo, 20 épocas no ImageNet) com pouca sobrecarga de treinamento. Além disso, para melhorar o comportamento condicional do portão durante o treinamento, introduzimos uma nova generalização da função de perda de moldagem de lote. Mostramos que nosso módulo de portão é capaz de aprender semântica significativa, apesar de operar localmente no nível de patch grosseiro. Validamos o MSViT nas tarefas de classificação e segmentação, onde ele leva a uma melhor relação de compensação entre precisão e complexidade.

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.