Популярный ИИ-исследователь Янник Кильчер (Yannic Kilcher) в своём детальном разборе анализирует научную статью, представляющую архитектуру Nyströmformer — новый алгоритм аппроксимации механизма Self-Attention. Авторы работы предложили элегантное решение фундаментальной проблемы квадратичной сложности классических трансформеров, основанное на математическом методе Найстрема. Ведущий подробно разбирает структуру алгоритма, его скрытые математические компромиссы, иронизирует над некоторыми формулировками авторов и оценивает впечатляющие результаты тестирования модели.
🧠 Квадратичное узкое горлышко классического Self-Attention 2:00
В основе любого современного трансформера лежит механизм самовнимания (Self-Attention). Как объясняет Янник Кильчер, в то время как Cross-Attention передаёт информацию между энкодером и декодером, Self-Attention работает внутри одного слоя, преобразуя входную последовательность в выходную последовательность равной длины. Задача этого слоя — перераспределить информацию так, чтобы каждый токен «знал» обо всех остальных релевантных элементах контекста.
Для этого за счёт линейных проекций создаются три матрицы: запросы (Queries, $Q$), ключи (Keys, $K$) и значения (Values, $V$). Ключевая и самая затратная операция — это перемножение матриц запросов и ключей ($Q \times K^T$), где каждый элемент выражает степень соответствия между запросом одного токена и ключом другого.
По мнению ведущего, именно здесь кроется главное аппаратное ограничение архитектуры:
- Матрица запросов имеет размерность $N \times D$, где $N$ — длина последовательности (количество токенов), а $D$ — внутренняя скрытая размерность.
- Матрица ключей имеет аналогичную размерность $N \times D$.
- Их произведение порождает гигантскую квадратную матрицу весов внимания размером $N \times N$, что создаёт квадратичную сложность по времени и памяти — $O(N^2)$.
Ситуация усложняется обязательным применением операции Softmax по строкам для получения корректного распределения вероятностей (гистограммы весов). По словам Кильчера, если бы не Softmax, исследователи могли бы легко использовать стандартные методы линейной алгебры для декомпозиции матрицы, поскольку её реальный ранг ограничен меньшим числом $D$ ($D \ll N$). Однако нелинейный Softmax требует вычисления всей строки целиком для её нормализации, что блокирует прямую факторизацию и заставляет разработчиков искать обходные пути.
🧮 Метод Найстрема: магия аппроксимации матриц 11:19
Авторы исследуемой статьи предлагают решить проблему вычисления полной матрицы с помощью классического метода Найстрема (Nyström method). Янник Кильчер в шутку переименовывает модель в «Найстример» (Nystreamer). Идея метода заключается в том, что любую большую матрицу можно мысленно разделить на четыре сектора (блока):
$$ \begin{pmatrix} A & B \ F & C \end{pmatrix} $$
Где блок $A$ — это небольшая квадратная матрица в левом верхнем углу, $B$ и $F$ — прямоугольные боковые блоки, а $C$ — огромный остаточный блок, занимающий львиную долю пространства.
Суть аппроксимации Найстрема, как демонстрирует ведущий, заключается в полном отказе от вычисления и хранения гигантского блока $C$. Его значения можно приближённо реконструировать по формуле:
$$C \approx F \cdot A^{-1} \cdot B$$
Вместо оригинальной сложности $N \times N$ мы переходим к работе со значительно меньшими матрицами размерностей $N \times M$ и $M \times M$, где $M$ — количество выбранных опорных точек, называемых ландмарками (landmarks).
🛠️ Главный математический «хак» Nyströmformer и перестановка операторов 15:40
Основная сложность интеграции метода Найстрема в трансформеры — это всё тот же Softmax. В наивном сценарии нам сначала пришлось бы рассчитать полную матрицу $N \times N$, применить к ней Softmax и только потом выделять из неё подматрицы $A, B, F$ для аппроксимации, что полностью лишает метод вычислительного смысла.
Чтобы обойти это препятствие, авторы Nyströmformer пошли на смелый шаг: они решили поменять порядок операций местами. По мнению Янника Кильчера, этот трюк математически не совсем корректен, так как Softmax и сэмплирование подматриц не являются коммутирующими операторами. Разработчики рассчитывают произведения для мелких подматриц, применяют Softmax к каждой из них независимо, а затем перемножают их через псевдообратную матрицу.
Ведущий с иронией разбирает математическое доказательство (лемму) сходимости алгоритма, приведённое в статье:
- В тексте утверждается, что аппроксимация сходится к истинному значению Self-Attention, если ландмарки (опорные точки) полностью совпадают с оригинальными матрицами запросов и ключей.
- Янник соглашается с этим фактом, комично отмечая: «Если вы выберете в качестве ландмарков абсолютно все ключи и запросы, аппроксимация действительно будет идеальной, потому что она перестанет быть аппроксимацией».
Тем не менее Кильчер вступается за авторов, указывая, что в дополнительном аппендиксе на GitHub (который было на удивление трудно найти) приведены реальные строгие доказательства ограничения погрешности в рамках $L_\infty$-нормы. Математика там рабочая, и за счёт того, что Softmax никогда не превышает единицу, ошибка аппроксимации остаётся в разумных пределах. Также Янник замечает техническую опечатку в формуле вычисления псевдоинверсии в основном тексте статьи, где пропущен знак минус первой степени (инверсии).
📦 Интуитивные аналоги Янника: полки с коробками и Word2Vec 23:07
Для объяснения того, почему метод Найстрема вообще способен адекватно предсказывать связи токенов, Янник приводит наглядную бытовую аналогию. Представьте, что вам нужно поднять коробку на определённую полку (оценить прямое взаимодействие запроса $Q_5$ и ключа $K_4$), но у вас нет прямой возможности это сделать. Вместо этого вы можете узнать, как $Q_5$ взаимодействует с опорным ключом $K_1$, затем как этот $K_1$ соотносится с опорным запросом $Q_1$, и наконец, как $Q_1$ взаимодействует с целевым ключом $K_4$. Из-за пересечения маршрутов в начале пути мы делим результат на взаимодействие базовых элементов $Q_1$ и $K_1$, получая достаточно точную оценку обходным путём.
Вторая интуитивная догадка ведущего касается легитимности локального применения Softmax к усечённым матрицам. Кильчер сравнивает это с техникой негативного сэмплирования (Negative Sampling) в классическом алгоритме Word2Vec. При обучении эмбеддингов слов невозможно рассчитать честный знаменатель Softmax по всему словарю, состоящему из миллионов слов. Поэтому его заменяют случайной выборкой негативных примеров, и эта аппроксимация отлично работает на практике. В Nyströmformer в качестве ландмарков берутся не просто случайные токены, а средние значения сегментов матриц (Segment Means), что делает такую выборку ещё более репрезентативной для нормализации распределения.
📊 Результаты бенчмарков: победа в Long Range Arena 43:19
Благодаря замене $O(N^2)$ вычислений на работу с ландмарками, сложность алгоритма падает до линейной — $O(N)$. На практике это даёт колоссальный выигрыш при масштабировании контекста, что Янник подтверждает конкретными цифрами из исследования:
| Длина последовательности (Sequence Length) | Память классического Transformer | Память Nyströmformer (64 ландмарка) |
|---|---|---|
| 512 токенов | 54 МБ | 35 МБ |
| 8000 токенов | 10 ГБ | 300 МБ |
Помимо экономии памяти, алгоритм обеспечивает мощный прирост скорости вычислений. Главным же сюрпризом для Янника стал график производительности модели в бенчмарке Long Range Arena (LRA), который один из авторов опубликовал в Twitter. Этот тест проверяет способность моделей обрабатывать сверхдлинные контексты и удерживать далёкие взаимосвязи в тексте.
Nyströmformer показал выдающиеся результаты:
- Модель смогла полностью повторить точность стандартного полноразмерного трансформера.
- Она значительно превзошла другие популярные линейные аппроксимации: Performer (использующий случайные признаки Фурье), Linformer (проектирующий длину последовательности) и Reformer (основанный на локально-чувствительном хэшировании).
Ведущий призывает относиться к экспериментам из одиночных статей с долей скептицизма, замечая, что ось ординат на графике LRA была смещена ради создания более драматичного эффекта. Тем не менее он признаёт результаты весьма многообещающими и заявляет, что этот график и математические выкладки из аппендикса обязаны были оказаться на первых страницах основного документа. Код проекта открыт, очень прост в реализации и доступен по ссылке в описании видео.
