Atlas: Многоуровневое внимание улучшает моделирование изображений с длинным контекстом

Atlas: Multi-Scale Attention Improves Long Context Image Modeling

March 16, 2025

Авторы: Kumar Krishna Agrawal, Long Lian, Longchao Liu, Natalia Harguindeguy, Boyi Li, Alexander Bick, Maggie Chung, Trevor Darrell, Adam Yala

cs.AI

Аннотация

Эффективное моделирование изображений большого размера остается давней проблемой в машинном обучении. Для решения этой задачи мы представляем Multi-Scale Attention (MSA). MSA основывается на двух ключевых идеях: (i) многоуровневые представления и (ii) двунаправленное взаимодействие между уровнями. MSA создает O(log N) уровней для представления изображения с постепенно более грубыми признаками и использует кросс-внимание для передачи информации между уровнями. Затем мы представляем Atlas — новую архитектуру нейронной сети, основанную на MSA. Мы демонстрируем, что Atlas значительно улучшает баланс между вычислительной сложностью и производительностью при моделировании изображений с длинным контекстом в высокоразрешенной версии ImageNet 100. При разрешении 1024px Atlas-B достигает точности 91.04%, что сопоставимо с ConvNext-B (91.92%), при этом работая в 4.3 раза быстрее. Atlas работает в 2.95 раза быстрее и на 7.38% точнее, чем FasterViT, а также в 2.25 раза быстрее и на 4.96% точнее, чем LongViT. В сравнении с MambaVision-S, Atlas-S демонстрирует на 5%, 16% и 32% более высокую точность при разрешениях 1024px, 2048px и 4096px соответственно, при схожем времени выполнения. Код для воспроизведения экспериментов и предобученные модели доступны по адресу 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.