Как Nyströmformer решает проблему квадратичной сложности в архитектуре Transformer

Yannic Kilcher 17,6 тыс. 48 мин 6 мин 11.02.2021
Главное

В новом видео популярный ИТ-исследователь и блогер Янник Килчер (Yannic Kilcher) подробно разбирает научную статью, посвященную алгоритму Nyströmformer. Этот метод предлагает эффективную аппроксимацию механизма самовнимания (Self-Attention) в архитектуре Transformer на основе метода Нистрёма. Главная цель разработки — преодолеть критическое «узкое горлышко» классических моделей, связанное с квадратичным ростом вычислительной сложности при увеличении длины обрабатываемых последовательностей.

🛑 Проблема квадратичной сложности в трансформерах 2:00

Классический механизм самовнимания (Self-Attention) лежит в основе успеха современных трансформеров, однако он страдает от серьезного ограничения: квадратичной сложности по времени и памяти $O(N^2)$. В одном слое внимания входная последовательность $X$ проецируется с помощью линейных матриц в три составляющие: запросы (Queries, $Q$), ключи (Keys, $K$) и значения (Values, $V$).

Процесс вычисления выглядит следующим образом:

Если размерность скрытых слоев ($D$) обычно меньше длины последовательности, то с математической точки зрения получившаяся матрица имеет низкий ранг (максимум $D$). Казалось бы, её можно легко декомпозировать на меньшие составляющие с помощью классических методов линейной алгебры. Однако на пути исследователей встает нелинейная операция softmax, которая применяется построчно для нормализации весов внимания и превращения их в валидное вероятностное распределение.

Из-за того, что softmax требует вычисления всей строки целиком для её корректного масштабирования, стандартные трюки декомпозиции матриц не работают напрямую. Попытки обучить сети вообще без нормализации, как отмечает Янник Килчер, пока показывают неудовлетворительные эмпирические результаты.

🧩 Суть метода Нистрёма и обход ограничений 11:19

Авторы исследуемой статьи предлагают использовать для решения этой проблемы метод Нистрёма (Nyström method). В линейной алгебре этот подход позволяет аппроксимировать большую матрицу, разбив её на блоки. Представим матрицу, разделенную на четыре сектора: верхний левый блок $A$ (небольшого размера), вытянутые боковые блоки $B$ и $F$, а также огромный остаточный блок $C$. Метод Нистрёма позволяет полностью отказаться от вычисления и хранения блока $C$, приближенно восстанавливая его по формуле:

$$C \approx F \times A^{-1} \times B$$

При таком подходе сложность вычислений падает, так как алгоритм оперирует матрицами меньшего размера ($N \times M$ и $M \times M$, где $M$ — количество выбранных опорных элементов).

Однако здесь возникает серьезный методологический конфликт с операцией нормализации. В наивном сценарии для аппроксимации матрицы весов внимания сначала нужно было бы рассчитать полную матрицу $N \times N$, применить к ней softmax и только потом выделять подматрицы для метода Нистрёма. Очевидно, что это полностью уничтожает весь вычислительный выигрыш.

Создатели Nyströmformer пошли на радикальный шаг: они фактически поменяли математический порядок операций местами. Вместо применения softmax к результату аппроксимации, они сначала вычисляют softmax для отдельных независимых подматриц меньшего размера, а затем перемножают их по формуле Нистрёма. По мнению Янника Килчера, с точки зрения строгой математики этот переход не является валидным, поскольку эти операторы не являются взаимно заменяемыми. Тем не менее, авторы приводят теоретические доказательства сходимости такого приближения к истинной матрице самовнимания.

💡 Интуиция: как это работает на пальцах 18:09

Чтобы объяснить, почему метод Нистрёма вообще работает в контексте нейросетей, Янник Килчер предлагает опереться на концепцию низкого ранга матрицы. Значения внутри матрицы внимания не являются абсолютно независимыми друг от друга. Это означает, что для описания всей системы требуется гораздо меньше физической информации, чем содержится в полной сетке $N \times N$ параметров.

Для вычислений алгоритм выбирает так называемые «ориентиры» (landmarks) — подмножество запросов и ключей, которые берут на себя роль репрезентативных точек всей последовательности.

Ведущий приводит наглядную последовательность для понимания того, как через эти ориентиры восстанавливается связь между любыми случайными элементами (например, запросом №5 и ключом №4):

  1. Нам неизвестно прямое взаимодействие между запросом 5 и ключом 4, и мы не хотим тратить ресурсы на его вычисление.
  2. Но мы знаем, как запрос 5 взаимодействует с опорным ключом 1 (это значение у нас есть).
  3. Также мы знаем, как опорный запрос 1 взаимодействует с ключом 4.
  4. Связав эти два шага через внутреннее взаимодействие самих ориентиров (опорного запроса 1 и опорного ключа 1), мы получаем обходной путь для оценки связи исходных элементов.

Для иллюстрации Килчер использует бытовой пример с перестановкой коробки: вместо того чтобы пытаться поднять тяжелую коробку напрямую на верхнюю полку, вы можете оценить необходимые усилия, если сначала поднимете её на одну промежуточную полку, затем на вторую и так далее. Да, возникнет небольшая погрешность из-за того, что каждый раз коробку приходится ставить и поднимать заново, но в целом вы получите довольно точное представление о финальной нагрузке. Математически этот обходной путь требует вычисления псевдообратной матрицы для взаимодействия самих ориентиров, что исключает проблемы с потенциальной необратимостью систем уравнений.

🛠️ Алгоритм Nyströmformer и архитектурные трюки 28:15

Вместо случайного выбора подмножества элементов для роли ориентиров, создатели Nyströmformer применили более изящное решение — усреднение сегментов (segment means). Алгоритм разбивает запросы и ключи на блоки и вычисляет их средние значения, что позволяет точнее сэмплировать распределение данных.

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

Для стабилизации процесса обучения авторы также добавили в архитектуру сквозную связь (skip connection). По мнению ведущего, этот элемент сильно напоминает идеи, заложенные в сетях LambdaNetworks, и помогает заметно ускорить сходимость модели.

Для обоснования спорного переноса softmax внутрь блоков Янник Килчер предлагает параллель из области дистрибутивной семантики — метод негативного сэмплирования (negative sampling) в моделях Word2Vec. При обучении векторных представлений слов невозможно честно рассчитать знаменатель для нормализации по всему гигантскому словарю. Вместо этого на практике берется случайная выборка «негативных» примеров, которая выступает в роли аппроксимации полной языковой среды. Килчер предполагает, что в Nyströmformer происходит аналогичный процесс: нормализация по усредненным ориентирам дает репрезентативное приближение для всего объема скрытых состояний.

📊 Эксперименты и победа в Long Range Arena 43:19

Благодаря замене полного самовнимания на блочные вычисления, теоретическая сложность Nyströmformer падает до строго линейной — $O(N)$. Внедрение этого метода дает колоссальный выигрыш по памяти при масштабировании контекста.

В видео приводятся следующие факты и замеры эффективности:

Самым интригующим результатом работы, который изначально отсутствовал в публикации, но был выложен авторами в Twitter, стали тесты на бенчмарке Long Range Arena (LRA). Данный набор задач спроектирован специально для проверки способности моделей удерживать сверхдлинные контекстные зависимости. Nyströmformer сумел не просто превзойти по качеству другие популярные «линейные» альтернативы — Performer (использующий случайные признаки Фурье), Linformer (работающий через проекцию длины последовательности) и Reformer (основанный на локально-чувствительном хешировании), — но и вплотную приблизился к показателям точности стандартного оригинального трансформера.

В завершение анализа Янник Килчер призывает традиционно относиться к любым академическим тестам с долей скепсиса. Он с иронией отмечает одну из лемм в приложении к статье, где авторы доказывают, что их аппроксимация идеально сходится с оригиналом, если в качестве ориентиров выбрать... абсолютно все исходные точки последовательности, что делает саму аппроксимацию бессмысленной. Тем не менее, математические границы погрешностей (L-infinity bounds), найденные в официальном репозитории на GitHub, подтверждают стабильность алгоритма, что делает Nyströmformer весьма перспективной технологией.

💬 Цитаты

«Значения внутри матрицы внимания не являются абсолютно независимыми друг от друга.»

Янник Килчер 18:23

«Вместо применения softmax к результату аппроксимации, они сначала вычисляют softmax для отдельных независимых подматриц меньшего размера»

Янник Килчер 17:03
👥 Спикер
🔗 Упомянутые сайты и проекты
📖 Термины
Self-Attention (Самовнимание)
Механизм в нейросетях, позволяющий элементам последовательности динамически определять степень важности других элементов для формирования собственного контекста.
Метод Нистрёма
Математический подход для аппроксимации больших матриц низкого ранга с помощью выборки опорных строк и столбцов.
Негативное сэмплирование (Negative Sampling)
Метод оптимизации, заменяющий ресурсоемкое вычисление нормализации по всему словарю на расчет по случайному подмножеству ложных примеров.
📊 Цифры
⚖️ Другая сторона
Искусственный интеллект Nyströmformer Self-Attention Yannic Kilcher Transformer