Atlas : L'attention multi-échelle améliore la modélisation d'images à contexte long

Atlas: Multi-Scale Attention Improves Long Context Image Modeling

March 16, 2025

Auteurs: Kumar Krishna Agrawal, Long Lian, Longchao Liu, Natalia Harguindeguy, Boyi Li, Alexander Bick, Maggie Chung, Trevor Darrell, Adam Yala

cs.AI

Résumé

La modélisation efficace d'images massives constitue un défi de longue date en apprentissage automatique. Pour y répondre, nous introduisons l'attention multi-échelle (Multi-Scale Attention, MSA). MSA repose sur deux idées clés : (i) les représentations multi-échelles et (ii) la communication bidirectionnelle entre échelles. MSA crée O(log N) échelles pour représenter l'image à travers des caractéristiques progressivement plus grossières et exploite l'attention croisée pour propager l'information entre les échelles. Nous présentons ensuite Atlas, une nouvelle architecture de réseau neuronal basée sur MSA. Nous démontrons qu'Atlas améliore significativement le compromis entre calcul et performance pour la modélisation d'images à contexte long dans une variante haute résolution d'ImageNet 100. À une résolution de 1024px, Atlas-B atteint une précision de 91,04 %, comparable à ConvNext-B (91,92 %) tout en étant 4,3 fois plus rapide. Atlas est 2,95 fois plus rapide et 7,38 % meilleur que FasterViT, et 2,25 fois plus rapide et 4,96 % meilleur que LongViT. En comparaison avec MambaVision-S, Atlas-S atteint respectivement 5 %, 16 % et 32 % de précision en plus à 1024px, 2048px et 4096px, tout en obtenant des temps d'exécution similaires. Le code pour reproduire nos expériences et les modèles pré-entraînés est disponible à l'adresse https://github.com/yalalab/atlas.

English

Efficiently modeling massive images is a long-standing challenge in machine learning. To this end, we introduce Multi-Scale Attention (MSA). MSA relies on two key ideas, (i) multi-scale representations (ii) bi-directional cross-scale communication. MSA creates O(log N) scales to represent the image across progressively coarser features and leverages cross-attention to propagate information across scales. We then introduce Atlas, a novel neural network architecture based on MSA. We demonstrate that Atlas significantly improves the compute-performance tradeoff of long-context image modeling in a high-resolution variant of ImageNet 100. At 1024px resolution, Atlas-B achieves 91.04% accuracy, comparable to ConvNext-B (91.92%) while being 4.3x faster. Atlas is 2.95x faster and 7.38% better than FasterViT, 2.25x faster and 4.96% better than LongViT. In comparisons against MambaVision-S, we find Atlas-S achieves 5%, 16% and 32% higher accuracy at 1024px, 2048px and 4096px respectively, while obtaining similar runtimes. Code for reproducing our experiments and pretrained models is available at https://github.com/yalalab/atlas.