Atlas: Atenção Multiescala Melhora a Modelagem de Imagens com Contexto Longo

Atlas: Multi-Scale Attention Improves Long Context Image Modeling

March 16, 2025

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

cs.AI

Resumo

Modelar imagens massivas de forma eficiente é um desafio de longa data no aprendizado de máquina. Para isso, introduzimos a Atenção Multi-Escala (MSA, na sigla em inglês). A MSA se baseia em duas ideias principais: (i) representações multi-escala e (ii) comunicação bidirecional entre escalas. A MSA cria O(log N) escalas para representar a imagem em características progressivamente mais grosseiras e aproveita a atenção cruzada para propagar informações entre as escalas. Em seguida, apresentamos o Atlas, uma nova arquitetura de rede neural baseada na MSA. Demonstramos que o Atlas melhora significativamente a relação entre desempenho e custo computacional na modelagem de imagens de contexto longo em uma variante de alta resolução do ImageNet 100. Na resolução de 1024px, o Atlas-B alcança 91,04% de precisão, comparável ao ConvNext-B (91,92%), enquanto é 4,3x mais rápido. O Atlas é 2,95x mais rápido e 7,38% melhor que o FasterViT, e 2,25x mais rápido e 4,96% melhor que o LongViT. Em comparações com o MambaVision-S, observamos que o Atlas-S alcança precisões 5%, 16% e 32% maiores em 1024px, 2048px e 4096px, respectivamente, mantendo tempos de execução semelhantes. O código para reproduzir nossos experimentos e os modelos pré-treinados está disponível em 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.