Atlas: La atención multi-escala mejora el modelado de imágenes con contexto extenso

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

Resumen

Modelar imágenes masivas de manera eficiente es un desafío de larga data en el aprendizaje automático. Para abordar este problema, presentamos la Atención Multi-Escala (MSA, por sus siglas en inglés). MSA se basa en dos ideas clave: (i) representaciones multi-escala y (ii) comunicación bidireccional entre escalas. MSA crea O(log N) escalas para representar la imagen a través de características progresivamente más gruesas y aprovecha la atención cruzada para propagar información entre escalas. Luego, presentamos Atlas, una novedosa arquitectura de red neuronal basada en MSA. Demostramos que Atlas mejora significativamente la relación entre rendimiento y cómputo en el modelado de imágenes de contexto largo en una variante de alta resolución de ImageNet 100. A una resolución de 1024px, Atlas-B alcanza un 91.04% de precisión, comparable a ConvNext-B (91.92%) mientras es 4.3 veces más rápido. Atlas es 2.95 veces más rápido y un 7.38% mejor que FasterViT, y 2.25 veces más rápido y un 4.96% mejor que LongViT. En comparaciones contra MambaVision-S, encontramos que Atlas-S logra un 5%, 16% y 32% más de precisión a 1024px, 2048px y 4096px respectivamente, mientras obtiene tiempos de ejecución similares. El código para reproducir nuestros experimentos y los modelos preentrenados está disponible en 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.