MSViT: Tokenización Dinámica de Escalas Mixtas para Transformadores de Visión

MSViT: Dynamic Mixed-Scale Tokenization for Vision Transformers

July 5, 2023

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

cs.AI

Resumen

Los tokens de entrada en los Transformers de Visión (ViT) tienen poco significado semántico, ya que se definen como parches regulares de tamaño fijo de la imagen de entrada, independientemente de su contenido. Sin embargo, procesar áreas de fondo uniformes de una imagen no debería requerir tanto cómputo como áreas densas y desordenadas. Para abordar este problema, proponemos un esquema de tokenización dinámica de escala mixta para ViT, denominado MSViT. Nuestro método introduce un mecanismo de compuerta condicional que selecciona la escala óptima de tokens para cada región de la imagen, de modo que el número de tokens se determina dinámicamente por entrada. El módulo de compuerta propuesto es ligero, independiente de la elección del backbone del transformer, y se entrena en pocas épocas (por ejemplo, 20 épocas en ImageNet) con un mínimo sobrecosto de entrenamiento. Además, para mejorar el comportamiento condicional de la compuerta durante el entrenamiento, introducimos una generalización novedosa de la función de pérdida de conformación por lotes (batch-shaping loss). Demostramos que nuestro módulo de compuerta es capaz de aprender semántica significativa a pesar de operar localmente a nivel de parches gruesos. Validamos MSViT en tareas de clasificación y segmentación, donde logra un mejor equilibrio entre precisión y complejidad.

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.