Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention
В статье разработана новая архитектура Infini-attention - она предполагает увеличение контекстного окна для механизма Linear Attention благодаря итерационному накоплению векторов KV в памяти M и применению delta rule к обновлению M.
Source: Arxive, YouTube YannicKilcher
Сначала напомню что такое Linear Attention - это механизм, который ускоряет вычисление матрицы attention (классическая вычисляется за O(N^2)) с помощью замены softmax произведения векторов QK на нелинейные функции (например сигмоиду - грубый пример ниже)
Основная структура алгоритма в рассматриваемой работе содержит два блока - Linear Attention и классический Multi-head Attention (PE - Position Embeddings).
Основной выход алгоритма это взвешенное сложение (параметр бета является обучаемым) результатов двух механизмов - классического Multi-Head Attention (A_dot) и Linear Attention с модификацией сжатия памяти(A_mem):
Напомню результат работы Multi-Head Attention - стандартный softmax от матрицы attention с нормировкой на размерность матрицы K и помноженный на вектор V:
Выход из Linear Attention представляется как произведение любой нелинейной функции от вектора Q на матрицу памяти M:
На замену нормализации корня из d приходит новое условие z. Матрица памяти и вектор z обновляются итерационно следующим образом:
Как раз в этой формуле мы видим накопление векторов K и V (само их произведение называется associative binding operator) в матрице M.
Однако авторы предлагают улучшение данного подхода: они добавляют проверку накопленных значений V с помощью delta rule - сначала извлекаются предыдущие значения вектора V путем умножения K на предыдущую матрицу M (числитель дроби) и вычитаются из нового:
Извлечение основано на ортогональности вектора K - если вектор помножится на самого себя то результатом будет 1, а если на другой вектор, то 0 (это возможно из за высокой размерности пространства):
То есть по сути здесь происходит проверка сохраненной памяти - если такие вектора уже были сохранен в памяти M, то получится вектор V и разность в скобках будет равна нулю, а значит обновления матрицы не произойдет, что сэкономит вычислительные ресурсы.
Финал - схема нового механизма Infini-Transformer, в которой показано сравнение с технологией Transformer-XL:
Контекст в Infini-Transformer проходит по всем блокам структуры, влияя на каждый трансформер, при этом на каждой итерации происходит накопление векторов KV в памяти M.
Безусловно механизм Infini-Transformer является абсолютно новым подходом в работе с большим контекстом. Важно, что при возможности сохранять и извлекать состояния векторов не требуется лишняя вычислительная мощность - все работает на перемножении матриц и векторов, хотя проблема O(N^2) в расчете выходов классического Attention остается, поэтому явного улучшения в скорости работы здесь ожидать не приходится. Я бы расценивал этот трансформер как дополнительную настройку классического Attention, способную улучшить результаты.
В статье представлено много результатов сравнения работы Infini-Transformer с другими технологиями в области оптимизации памяти и работы с длинным контекстом, однако самым интересным результатом для меня является результаты теста суммаризации теста BookSum
