Быстрее EfficientNet в 4.5 раза: Янник Килчер объясняет устройство LambdaLayers

Yannic Kilcher 48,9 тыс. 59 мин 7 мин 17.10.2020
Главное

Известный ИИ-исследователь и блогер Янник Килчер (Yannic Kilcher) представил подробный разбор научной статьи «LambdaNetworks: Modeling long-range Interactions without Attention». В этой работе анонимные авторы предложили альтернативу классическому механизму внимания (attention), способную значительно ускорить обучение нейросетей при сохранении высокой точности. Предложенная архитектура Lambda ResNets обходит по эффективности популярные модели EfficientNet и стандартные ResNet на датасете ImageNet, предлагая свежий взгляд на проблему долгосрочных зависимостей в данных.

🚀 Прорыв на ImageNet: Проблема трансформеров и альтернатива от LambdaNetworks 0:00

В мире машинного обучения зафиксирован очередной прорыв на датасете ImageNet. На этот раз превосходные результаты показала архитектура под названием Lambda ResNets. Как утверждает ведущий Янник Килчер, данная модель превосходит EfficientNet и классические ResNet не только по точности топ-1, но и по соотношению качества работы и времени, затраченного на обучение. Согласно данным из исследуемой статьи, Lambda ResNets работают примерно в 4.5 раза быстрее, чем EfficientNets, существенно улучшая компромисс между скоростью и точностью классификации изображений на разных масштабах. В последнее время в компьютерном зрении доминировали трансформеры, однако их применение требовало либо сильного даунсэмплинга (например, разбиения изображения на патчи размером $16 \times 16$), либо огромных объемов данных и вычислительных мощностей. Авторы статьи «LambdaNetworks: Modeling long-range Interactions without Attention», находящейся на рецензировании в рамках конференции ICLR 2021, обещают гораздо более эффективное решение. Хотя работа опубликована анонимно, Янник Килчер решает сосредоточиться исключительно на ее содержании, признавая, что глубоко разобраться в предложенном математическом аппарате было непросто.

🧠 Суть Lambda-слоёв: Линейные функции вместо карт внимания 2:11

Разработчики представляют Lambda-слой как общий каркас для улавливания дальнодействующих связей между входными данными и структурированным контекстом — например, центральным пикселем и окружающим его фоном. Метод работает за счет трансформации доступного контекста в линейные функции, называемые «лямбдами» (термин отсылает к лямбда-исчислению), которые затем применяются к каждому входу отдельно. По словам Янника Килчера, авторы заявляют о возможности глобального, локального или маскированного моделирования контентных и позиционных взаимодействий. Поскольку архитектура обходится без вычисления дорогостоящих карт внимания (attention maps), её можно применять к последовательностям длиной в тысячи элементов и к изображениям высокого разрешения. В трансформерах ключевым ограничивающим фактором выступает память, объем которой невозможно увеличить простым ожиданием окончания вычислений. Однако Янник Килчер критически оценивает заявления авторов о легкой работе с высоким разрешением: по его мнению, на практике это неосуществимо без использования так называемого локального внимания (local attention).

Для понимания масштаба проблемы ведущий приводит простой пример:

Чтобы обойти это ограничение, исследователи используют локальное внимание, которое Янник Килчер метафорически называет «динамической сверткой». В обычной свертке ядро зафиксировано, а локальное внимание формирует динамическое ядро вокруг целевого пикселя. Lambda-слои действуют схожим образом, ограничивая поле зрения локальной областью, но принципиальное отличие кроется в ином: они сначала абстрагируют и суммируют контекст в матрицу малой размерности без оглядки на сам запрос (query). Полученная линейная функция затем умножается на вектор запроса, устраняя нелинейность в виде софтмакса (softmax) над попарными связями, свойственную стандартному Attention.

📊 Математический разбор: Чем Lambda-слой отличается от классического Attention 17:19

Математическая структура Lambda-слоев довольно сложна для восприятия. На вход подаются матрицы элементов $X$ размера $N \times D$ (где $N$ — размер входа, а $D$ — количество каналов) и контекст $C$ размера $M \times D$. Зачастую $X$ и $C$ представляют собой один и тот же объект. Янник Килчер напоминает, как устроен классический расчет внимания:

  1. Входные данные умножаются на обучаемые матрицы весов $W_Q$, $W_K$ и $W_V$, порождая матрицы запросов $Q$, ключей $K$ и значений $V$.
  2. Запросы $Q$ и ключи $K$ перемножаются, формируя карту внимания размера $N \times M$.
  3. К этой карте применяется операция Softmax для нормализации, после чего результат умножается на значения $V$, выдавая итоговый выход $Y$ размера $N \times D$.

В LambdaNetworks процедура трансформации контекста кардинально меняется. На первом этапе вычисляются матрицы $Q$ ($N \times K$), $K$ ($M \times K$) и $V$ ($M \times D$). Главная хитрость, по мнению Янника Килчера, заключается в том, что операция Softmax применяется исключительно к матрице ключей $K$ вдоль размерности контекста $M$. В результате генерируется фиксированное количество ($K$) карт внимания над $M$ входами. Затем нормализованные ключи перемножаются со значениями $V$, и происходит свертывание (контракция) по размерности $M$. На выходе формируется матрица контентной лямбды $\lambda_c$ размера $K \times D$, из которой полностью исчезла опасная переменная $M$. Контекст оказывается полностью заархивирован до того, как он получит шанс провзаимодействовать с запросами $N$. Финальный шаг тривиален: матрица запросов $Q$ просто перемножается с $\lambda_c$, возвращая выходной тензор $Y$. Килчер проводит аналогию с концепцией «линформеров» (Linformer) и архитектур с фиксированными обучаемыми запросами, где сеть на каждом слое создает устойчивую псевдопоследовательность длины $K$. Запросы верхнегольного уровня лишь выбирают нужную точку в этом промежуточном пространстве признаков.

📍 Ловушка позиционного кодирования: Почему «магия» работает 40:16

Как отмечает Янник Килчер, голая математика контентной суммаризации, скорее всего, не показала бы жизнеспособных результатов на практике. Чтобы заставить Lambda-слои работать, авторам пришлось внедрить мощный механизм позиционного кодирования. Они ввели фиксированные обучаемые параметры взаимодействия $E$ в виде тензора размером $N \times M \times K$. По мнению Килчера, появление трехмерного тензора, включающего размерности $N$ и $M$, концептуально разрушает первоначальный нарратив авторов об отказе от громоздких карт попарного взаимодействия. Разница лишь в том, что эта огромная матрица позиционных связей не рассчитывается динамически на лету из входных данных, а статически зашита в веса конкретного слоя. Это дает колоссальное преимущество при масштабировании: тензор позиционных эмбеддингов не размножается вместе с размером батча (batch size). При обработке пакета данных память под размерность $N \times M$ выделяется ровно один раз.

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

Янник Килчер делится критическим наблюдением: эксперименты в статье наглядно показывают, что именно позиционный путь обеспечивает львиную долю итогового качества работы нейросети. Если полностью убрать позиционные кодирования, модель перестает обучаться, тогда как при отключении контентной лямбды результаты остаются вполне приемлимыми. Таким образом, реальная сила LambdaNetworks, по мнению ведущего, заключается не в хитрой замене Attention линейной функцией, а в эффективной реализации фиксированных пространственных весов, которые можно представить как аналог классического смещения (bias) в линейном уравнении вида $y = wx + b$.

⏱️ Анализ производительности и элегантная реализация в коде 53:09

Сравнивая Lambda-слой со стандартным Attention, авторы разделяют затраты памяти на контентное и позиционное взаимодействия. Главный выигрыш заключается в отсутствии мультипликативного влияния размера батча на квадратичную размерность. С точки зрения временной сложности ($Time\ Complexity$), чуда не произошло: алгоритм по-прежнему выполняет операции порядка $N \times M$, однако выигрыш в скорости достигается за счет того, что данные операции можно свести к быстрым сверточным примитивам, хорошо оптимизированным на уровне современных библиотек. В тестах на ImageNet модель Lambda ResNet демонстрирует выдающуюся пропускную способность (throughput), особенно при малых значениях гиперпараметра $K$ (например, $K=16$ или $K=8$), резко опережая локальное самовнимание, которое сильно страдает от неэффективной реализации индексных сдвигов в современных фреймворках.

Архитектурные тесты выявили любопытные нюансы интеграции новых слоев в классический ResNet-50:

В завершение обзора Янник Килчер выражает восхищение неофициальной реализацией этого алгоритма от известного open-source разработчика Фила Ванга (Phil Wang), публикующегося под ником lucidrains. Автор кода сумел воссоздать архитектуру Lambda-слоя практически сразу после выхода препринта статьи. Программа написана невероятно лаконично благодаря использованию библиотеки einops, заменяющей громоздкие и малопонятные операции reshape, squeeze и unsqueeze на наглядные декларативные перестановки осей тензоров. Особый восторг у Янника вызвало то, что разработчик использовал в коде на Python настоящие греческие буквы для именования переменных, включая сам символ лямбды ($\lambda$), что добавило математической элегантности практической реализации.

💬 Цитаты

«Внимание переносится на первую часть и становится своего рода фиксированным.»

Янник Килчер 35:17

«Позиционные кодирования фактически важнее в экспериментах: если их убрать, ничего не работает.»

Янник Килчер 46:54
👥 Спикер
🔗 Упомянутые сайты и проекты
📖 Термины
Attention (Механизм внимания)
Метод динамической маршрутизации информации в нейросетях, вычисляющий попарные взаимодействия между всеми элементами входа.
Lambda layer (Лямбда-слой)
Альтернативный слой нейросети, который абстрагирует контекст в линейную функцию перед применением к запросам.
EfficientNet
Семейство сверточных нейросетей, оптимизированных по точности и вычислительной эффективности.
Einops
Библиотека для Python, позволяющая гибко и наглядно манипулировать размерностями тензоров.
📊 Цифры
⚖️ Другая сторона
Искусственный интеллект LambdaNetworks Yannic Kilcher Self-Attention ImageNet