Поиск по сайту:

Оценка необходимости механизмов Mamba в задачах визуального распознавания-MambaOut

Введение

Трансформаторы являются основой таких моделей, как BERT, серия GPT и ViT. Однако его механизм внимания имеет квадратичную сложность, что затрудняет работу с длинными последовательностями. Для решения этой проблемы были разработаны различные миксеры токенов с линейной сложностью.

В последнее время модели на основе RNN привлекли внимание благодаря их эффективному обучению и выводу длинных последовательностей и показали себя многообещающими в качестве основы для больших языковых моделей.

Вдохновленные этими возможностями, исследователи изучили возможность использования Mamba в задачах визуального распознавания, что привело к созданию таких моделей, как Vision Mamba, VMamba, LocalMamba и PlainMamba. Несмотря на это, эксперименты показывают, что модель пространства состояний или модели зрения на основе SSM уступают по эффективности современным сверточным моделям и моделям, основанным на внимании.

Эта недавняя статья не посвящена разработке новых визуальных моделей Mamba. Вместо этого исследуется важнейший исследовательский вопрос: необходима ли Мамба для задач визуального распознавания?

Что такое Мамба?

Mamba — это архитектура глубокого обучения, разработанная исследователями из Университета Карнеги-Меллон и Принстонского университета и предназначенная для устранения ограничений моделей трансформаторов, особенно для длинных последовательностей. Он использует модель последовательности структурированного пространства состояний (S4), сочетающую сильные стороны непрерывных, рекуррентных и сверточных моделей для эффективной обработки длинных зависимостей и данных с нерегулярной выборкой.

Недавно исследователи адаптировали Mamba для задач компьютерного зрения, аналогично тому, как используются Vision Transformers (ViT). Vision Mamba (ViM) повышает эффективность за счет использования двунаправленной модели пространства состояний (SSM), отвечающей высоким вычислительным требованиям традиционных трансформаторов, особенно для изображений с высоким разрешением.

Мамба Архитектура

Mamba расширяет модель S4, вводя уникальный механизм выбора, который адаптирует параметры на основе входных данных, позволяя сосредоточиться на соответствующей информации в последовательностях. Эта изменяющаяся во времени структура повышает эффективность вычислений.

Mamba также использует аппаратно-ориентированный алгоритм для эффективных вычислений на современном оборудовании, таком как графические процессоры, оптимизируя производительность и использование памяти. Архитектура объединяет дизайн SSM с блоками MLP, что делает ее подходящей для различных типов данных, включая язык, аудио и геномику.

Варианты мамбы

  • MambaByte: языковая модель без токенов, которая обрабатывает необработанные последовательности байтов, устраняя токенизацию и связанные с ней предвзятости.
  • Mamba Mixture of Experts (MOE): интегрирует Mixture of Experts с Mamba, повышая эффективность и масштабируемость за счет чередования слоев Mamba и MOE.
  • Вижн Мамба (ВиМ): ViM адаптирует SSM для обработки визуальных данных, используя двунаправленные блоки Mamba для кодирования визуальных последовательностей. Это снижает вычислительные требования и повышает производительность при выполнении таких задач, как классификация ImageNet, обнаружение объектов COCO и семантическая сегментация ADE20k.
  • Jamba: Разработанный AI21 Labs, Jamba представляет собой гибридный преобразователь и архитектуру SSM Mamba с 52 миллиардами параметров и контекстным окном из 256 тысяч токенов.

Демо

Прежде чем мы начнем работать с моделью, мы клонируем репозиторий и установим несколько необходимых пакетов.

!pip install timm==0.6.11
!git clone https://github.com/yuweihao/MambaOut.git
!pip install gradio

Кроме того, мы добавили ссылку, которую можно использовать для доступа к блокноту, в котором выполняются шаги и выполняются выводы с помощью MambaOut.

cd /MambaOut

Ячейка ниже поможет вам запустить веб-приложение Gradio.

!python gradio_demo/app.py

RNN-подобные модели и причинное внимание

На иллюстрации ниже объясняется механизм причинного внимания и RNN-подобных моделей с точки зрения памяти, где _xi_ представляет входной токен на i-м шаге.

(a) Причинное внимание: сохраняет ключи (k) и значения (v) всех предыдущих токенов в памяти. Память обновляется путем постоянного добавления ключа и значения текущего токена, что делает его без потерь. Однако вычислительная сложность интеграции старой памяти с текущими токенами возрастает по мере удлинения последовательности. Таким образом, внимание хорошо работает с короткими последовательностями, но плохо работает с более длинными.

(b) RNN-подобные модели: сжимайте предыдущие токены в скрытое состояние фиксированного размера (h), которое служит памятью. Этот фиксированный размер означает, что память RNN по своей сути является памятью с потерями и не может соответствовать емкости памяти без потерь моделей внимания. Тем не менее, RNN-подобные модели превосходно справляются с обработкой длинных последовательностей, поскольку сложность объединения старой памяти с текущими входными данными остается постоянной, независимо от длины последовательности.

Мамба особенно хорошо подходит для задач, требующих смешивания причинных токенов из-за ее повторяющихся свойств. В частности, Мамба отлично справляется с задачами со следующими характеристиками:

1. Задача предполагает обработку длинных последовательностей. 2. Задача требует смешивания причинно-следственных токенов.

Возникает следующий вопрос: имеют ли задачи визуального распознавания очень длинные последовательности?

Для классификации изображений в ImageNet типичный размер входного изображения составляет 224x224, в результате чего получается 196 токенов с размером патча 16x16. Это число намного меньше пороговых значений для задач с длинной последовательностью, поэтому классификация ImageNet не считается задачей с длинной последовательностью.

Для обнаружения объектов и сегментации экземпляров на COCO с размером изображения 800x1280 и для семантической сегментации на ADE20K (ADE20K — широко используемый набор данных для задачи семантической сегментации, состоящий из 150 семантических категорий. Набор данных включает 20 000 изображений в обучении набор и 2000 изображений в наборе проверки) с размером изображения 512x2048, количество токенов около 4000 с размером патча 16x16. Поскольку 4000 токенов превышают порог для небольших последовательностей и близки к базовому порогу, как обнаружение COCO, так и сегментация ADE20K считаются задачами для длинных последовательностей.

Фреймворк MambaOut

Общая структура MambaOut

На рис. (a) представлена общая структура MambaOut для визуального распознавания: MambaOut предназначен для визуального распознавания и имеет иерархическую архитектуру, аналогичную ResNet. Он состоит из четырех этапов, каждый из которых имеет разные размеры канала, обозначаемые как _Di_. Эта иерархическая структура позволяет модели обрабатывать визуальную информацию на нескольких уровнях абстракции, повышая ее способность распознавать сложные закономерности в изображениях.

(b) Архитектура закрытого блока CNN: Блок Gated CNN является компонентом платформы MambaOut. Он отличается от блока Мамбы тем, что не включает в себя модель пространства состояний (SSM). Хотя оба блока используют сверточные нейронные сети (CNN) со стробирующими механизмами для регулирования потока информации, отсутствие SSM в блоке Gated CNN означает, что он не обладает такой же способностью обрабатывать длинные последовательности и временные зависимости, как блок Mamba, который включает в себя SSM. для этих целей.

Основное различие между Gated CNN и блоком Mamba заключается в наличии модели пространства состояний (SSM).

В MambaOut в качестве микшера токенов Gated CNN, аналогично ConvNeXt, используется глубокая свертка с размером ядра 7x7. Подобно ResNet, MambaOut построен с использованием четырехэтапной структуры путем объединения блоков Gated CNN на каждом этапе, как показано на рисунке.

Прежде чем двигаться дальше, выскажем гипотезу о необходимости внедрения Мамбы для визуального распознавания. Гипотеза 1: Нет необходимости внедрять SSM для классификации изображений в ImageNet, поскольку эта задача не соответствует характеристике 1 или характеристике 2. Гипотеза 2. По-прежнему имеет смысл продолжить изучение потенциала SSM для визуального обнаружения и сегментации, поскольку эти задачи соответствуют характеристике 1, несмотря на то, что они не соответствуют характеристике 2.

Обучение

Классификация изображений в ImageNet

  • ImageNet используется в качестве эталона для классификации изображений и содержит 1,3 миллиона обучающих изображений и 50 000 проверочных изображений.
  • Обучение проводится по схеме DeiT без дистилляции, включая различные методы увеличения данных и методы регуляризации.
  • Для обучения используется оптимизатор AdamW с правилом масштабирования скорости обучения lr=размер пакета/1024 * 10^-3, что приводит к скорости обучения 0,004 при размере пакета 4096.
  • Модели MambaOut реализованы с использованием библиотек PyTorch и timm и обучены на TPU v3.

Результаты классификации изображений на ImageNet

  • Модели MambaOut, которые не включают SSM, неизменно превосходят визуальные модели Mamba во всех размерах моделей на ImageNet.
  • Например, модель MambaOut-Small достигает точности высшего уровня 84,1 %, превосходя по производительности LocalVMamba-S на 0,4 %, но при этом требует только 79 % MAC-адресов.
  • Эти результаты подтверждают гипотезу 1, предполагающую, что введение SSM для классификации изображений в ImageNet не является необходимым.
  • Модели Visual Mamba в настоящее время значительно отстают от современных моделей свертки и внимания в ImageNet.
  • Например, CAFormer-M36 превосходит все визуальные модели Mamba сопоставимого размера по точности более чем на 1%.
  • Будущим исследованиям, направленным на опровержение гипотезы 1, возможно, потребуется разработать визуальные модели Mamba с миксерами токенов свертки и SSM для достижения современной производительности в ImageNet.

Обнаружение объектов и сегментация экземпляров на COCO

  • COCO 2017 используется в качестве эталона для обнаружения объектов и сегментации экземпляров.
  • MambaOut используется в качестве основы в Mask R-CNN и инициализируется с помощью весов, предварительно обученных в ImageNet.
  • Обучение проводится по стандартному графику 1× из 12 эпох, при этом размер обучающих изображений изменяется так, чтобы короткая сторона составляла 800 пикселей, а более длинная сторона не превышала 1333 пикселей.
  • Оптимизатор AdamW используется со скоростью обучения 0,0001 и общим размером пакета 16.
  • Реализация осуществляется с использованием библиотек PyTorch и mmdetection с точностью FP16 для экономии затрат на обучение.
  • Эксперименты проводятся на 4 графических процессорах NVIDIA 4090.

Результаты обнаружения объектов и сегментации экземпляров на COCO

  • Хотя MambaOut может превосходить некоторые визуальные модели Mamba в обнаружении объектов и сегментации экземпляров на COCO, он по-прежнему отстает от современных визуальных моделей Mamba, таких как VMamba и LocalVMamba.
  • Например, MambaOut-Tiny в качестве основы для Mask R-CNN отстает от VMamba-T на 1,4 APb и 1,1 APm.
  • Эта разница в производительности подчеркивает преимущества интеграции Mamba в визуальные задачи с длинными последовательностями, что подтверждает Гипотезу 2.
  • Однако визуальная Mamba по-прежнему демонстрирует значительный разрыв в производительности по сравнению с современными гибридными моделями свертки-внимания, такими как TransNeXt. Visual Mamba необходимо продемонстрировать свою эффективность, превзойдя другие современные модели в задачах визуального обнаружения.

Семантическая сегментация на ADE20K

  • ADE20K используется в качестве эталона для задачи семантической сегментации, включающей 150 семантических категорий с 20 000 изображений в обучающем наборе и 2000 изображений в проверочном наборе.
  • Mamba используется в качестве основы для UperNet с инициализацией из предварительно обученных весов ImageNet.
  • Обучение проводится с помощью оптимизатора AdamW со скоростью обучения 0,0001 и размером пакета 16 на 160 000 итераций.
  • Реализация осуществляется с использованием библиотек PyTorch и mmsegmentation, при этом эксперименты проводятся на четырех графических процессорах NVIDIA 4090 с использованием точности FP16 для повышения скорости обучения.

Результаты семантической сегментации на ADE20K

  • Подобно обнаружению объектов в COCO, тенденция производительности семантической сегментации в ADE20K показывает, что MambaOut может превосходить некоторые визуальные модели Mamba, но не может соответствовать результатам современных моделей Mamba.
  • Например, LocalVMamba-T превосходит MambaOut-Tiny на 0,5 млн единиц как в одномасштабных (SS), так и в многомасштабных (MS) оценках, что дополнительно подтверждает гипотезу 2 эмпирически.
  • Кроме того, визуальные модели Mamba по-прежнему демонстрируют заметный дефицит производительности по сравнению с более продвинутыми гибридными моделями, которые объединяют механизмы свертки и внимания, такими как SG-Former и TransNeXt.
  • Visual Mamba необходимо и дальше демонстрировать свои сильные стороны в моделировании длинных последовательностей, добившись более высокой производительности в задаче визуальной сегментации.

Заключение

Механизм Мамбы лучше всего подходит для задач с длинными последовательностями и авторегрессионными характеристиками. Mamba демонстрирует потенциал для задач визуального обнаружения и сегментации, которые соответствуют характеристикам длинной последовательности. Модели MambaOut, которые превосходят все визуальные модели Mamba в ImageNet, но все же отстают от современных визуальных моделей Mamba.

Однако из-за ограничений вычислительных ресурсов в этой статье основное внимание уделяется проверке концепции Мамбы для визуальных задач. Будущие исследования могут дополнительно изучить концепции Mamba и RNN, а также интеграцию RNN и Transformer для больших языковых моделей (LLM) и больших мультимодальных моделей (LMM), что потенциально приведет к новым достижениям в этих областях.

Ссылки

  • Оригинальная исследовательская статья: - MambaOut: действительно ли нам нужна мамба для зрения?
  • Vision Mamba: эффективное обучение визуальному представлению с помощью двунаправленной модели пространства состояний
  • Мамба (архитектура глубокого обучения)