BiFormer : Vision Transformer avec une attention de routage à deux niveaux

En tant que bloc de construction fondamental des transformateurs de vision, l’attention constitue un outil puissant pour capturer les dépendances à longue portée. Toutefois, cette puissance s’accompagne d’un coût élevé : elle engendre une charge computationnelle considérable et une empreinte mémoire importante, car elle implique le calcul d’interactions par paires entre tous les tokens à toutes les positions spatiales. Une série de travaux ont tenté de réduire ce problème en introduisant une sparsité artificielle et indépendante du contenu dans l’attention, par exemple en restreignant l’opération d’attention à des fenêtres locales, des bandes axiales ou des fenêtres dilatées. À la différence de ces approches, nous proposons une nouvelle attention sparse dynamique basée sur un routage à deux niveaux, permettant une allocation plus flexible des calculs, tout en étant sensible au contenu. Plus précisément, pour une requête donnée, les paires clé-valeur inappropriées sont d’abord filtrées à l’échelle grossière des régions, puis une attention fine à niveau token-to-token est appliquée dans l’union des régions candidates restantes (c’est-à-dire les régions routées). Nous proposons une implémentation simple mais efficace de l’attention à routage à deux niveaux, exploitant la sparsité pour réduire à la fois la charge computationnelle et l’empreinte mémoire, tout en n’impliquant que des multiplications matricielles denses compatibles avec les GPU. En s’appuyant sur cette attention à routage à deux niveaux, nous présentons un nouveau transformateur de vision généraliste, nommé BiFormer. Grâce à sa capacité à s’attacher à un petit sous-ensemble de tokens pertinents de manière adaptative à la requête, sans être perturbé par des éléments non pertinents, BiFormer atteint à la fois de bonnes performances et une haute efficacité computationnelle, particulièrement dans les tâches de prédiction dense. Des résultats expérimentaux sur plusieurs tâches de vision par ordinateur, telles que la classification d’images, la détection d’objets et la segmentation sémantique, confirment l’efficacité de notre conception. Le code est disponible à l’adresse suivante : \url{https://github.com/rayleizhu/BiFormer}.