Удалите Softmax Attention, и сложность уменьшится до O(n)

Как мы все знаем, хотя модель Transformer, основанная на механизме Attention, имеет хорошую параллельную производительность, ее пространственная и временная сложность невелики.$\mathcal{O}(n^2)$уровень,$n$длина последовательности, поэтому, когда$n$Когда он относительно велик, вычислительная нагрузка модели Transformer невыносима. В последнее время много работы было посвящено уменьшению вычислительной нагрузки модели Transformer, например, сокращению модели, квантованию, дистилляции и другим методам оптимизации, или модификации структуры Attention, чтобы ее сложность можно было уменьшить до$\mathcal{O}(nlog⁡n)$четное$\mathcal{O}(n)$

бумага«Трансформаторы - это RNN: быстрые авторегрессионные преобразователи с линейным вниманием»В нем упоминался метод линеаризации внимания (Linear Attention), который вызвал у меня интерес, а затем я прочитал несколько связанных блогов и дал хорошие результаты, и, наконец, обобщил свое понимание линеаризованного внимания в этой статье.

Attention

Наиболее популярным механизмом внимания являетсяScaled-Dot Attention,Сейчас

здесь$\boldsymbol{Q}\in \mathbb{R}^{n\times d_k}, \boldsymbol{K}\in \mathbb{R}^{m\times d_k}, \boldsymbol{V}\in \mathbb{R}^{m\times d_v}$, для простоты я не показывал коэффициент масштабирования Внимания$\frac{1}{\sqrt{d}}$. В этой статье нас в основном интересует сцена «Самовнимание», поэтому для удобства введения мы установили$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}\in \mathbb{R}^{n\times d}$

Удалить Софтмакс

Читатели могут не подумать, что ключевым фактором, ограничивающим производительность Attention, на самом деле является Softmax в определении! На самом деле, простой вывод приводит к такому выводу.$QK^T$На этом шаге мы получаем$n\times n$Матрица, тогда нам нужно сделать Softmax

к одному$1\times n$Softmax вектор-строка , временная сложность$O(n)$, но для$n\times n$Сделайте Softmax для каждой строки матрицы, и временная сложность$O(n^2)$

Если нет Софтмакса, то формула Внимания становится произведением трех матриц$\boldsymbol{QK^{\top}V}$, а умножение матриц удовлетворяет ассоциативной скорости, поэтому мы можем сначала вычислить$\boldsymbol{K^{\top}V}$, получить$d\times d$матрица (временная сложность этого шага$O(d^2n)$), затем используйте$Q$Умножьте его налево (временная сложность этого шага$O(d^2n)$),так как$d \ll n$, поэтому приблизительная временная сложность равна$O(n)$

Для базы BERT,$d=64$вместо 768, почему? Поскольку 768 на самом деле получается путем сращивания нескольких головок, и каждая головка$d=64$

То есть удаление сложности «Внимание» из Softmax может быть сведено к самому идеальному линейному уровню.$\mathcal{O}(n)$! Очевидно, это наша конечная цель: линейное внимание.

общее определение

Вопрос в том, можно ли считать удаление Softmax Внимание? Может ли он по-прежнему иметь стандартный эффект внимания? Чтобы ответить на этот вопрос, мы сначала перепишем определение Scaled-Dot Attention эквивалентным образом (все векторы в этой статье являются векторами-столбцами)

Вот небольшое объяснение, прежде всего мы знаем$\boldsymbol{Q},\boldsymbol{K}\in \mathbb{R}^{n\times d}$,сделать$\boldsymbol{M} = \boldsymbol{Q}\times \boldsymbol{K^{\top}}$, согласно правилу умножения матриц,$\boldsymbol{M}$Первая строка$\boldsymbol{Q}$Умножьте первую строку$\boldsymbol{K^{\top}}$Все столбцы полученного

$Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i$представляет первую из окончательной выходной матрицы$i$Ряд

$\boldsymbol{q}_i^{\top}$выражать$\boldsymbol{Q}\in \mathbb{R}^{n\times d}$первая из матрицы$i$строка (вектор-строка)

$\boldsymbol{k}_j$выражать$\boldsymbol{K^{\top}}\in \mathbb{R}^{d\times n}$первая из матрицы$j$столбец (вектор-столбец)

$\boldsymbol{v}_j$выражать$V^{\top}\in \mathbb{R}^{d\times n}$первая из матрицы$j$столбец (вектор-столбец)

Таким образом, Scaled-Dot Attention на самом деле основано на$e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}$для весовой пары$\boldsymbol{v}_j$Сделайте средневзвешенное значение. Таким образом, мы можем предложить обобщенное определение внимания.

то есть поставить$e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}$заменить$\boldsymbol{q}_i,\boldsymbol{k}_i$общая функция$\text{sim}(\boldsymbol{q}_i,\boldsymbol{k}_j)$, чтобы сохранить аналогичные характеристики распределения внимания, мы требуем$\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\geq 0$Хэн установил. То есть, если мы хотим определить новое внимание, мы должны сохранить форму уравнения (3) и удовлетворить$\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\geq 0$

Эта общая форма внимания также известна как нелокальная сеть в CV из статьи«Нелокальные нейронные сети»

несколько примеров

Если вы напрямую удалите Softmax, то это$\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j) = \boldsymbol{q}_i^{\top}\boldsymbol{k}_j$, проблема в том, что внутренний продукт не гарантирует неотрицательность, так что это пока не разумный выбор. Ниже мы приводим несколько возможных вариантов

Стоит отметить, что первые два вида линейного внимания, представленные ниже, относятся к полю CV, а третий — кСу ЦзяньлиньБольшой парень задумал (помимо следующего введения, есть ещеEMANetУлучшения внимания в области резюме и т. д.)

Форма функции ядра

Естественная мысль: если$\boldsymbol{q}_i, \boldsymbol{k}_j$Каждый элемент неотрицательный, поэтому скалярный продукт, естественно, неотрицательный. Для этого мы можем дать$\boldsymbol{q}_i, \boldsymbol{k}_j$Добавьте функцию активации к каждому$\phi,\varphi$,Сейчас

в$\phi(\cdot), \varphi(\cdot)$является функцией активации с неотрицательным диапазоном. Документы, упомянутые в начале статьиТрансформеры - это RNNвыбрано$\phi(x)=\varphi(x)=\text{elu}(x)+1$,в

Общий$\alpha$ценность$[0.1, 0.3]$

Если вам нужно рассказать историю, уравнение (4) можно связать с «ядерным методом», особенно$\phi=\varphi$час,$\phi$эквивалентна функции ядра, и$\langle \phi(\boldsymbol{q}_i), \phi(\boldsymbol{k}_j)\rangle$Это внутренний продукт, определяемый функцией ядра. Думая об этом, можно обратиться к статье«Рассечение трансформатора: единое понимание внимания трансформатора через призму ядра», не распространяйтесь слишком много здесь

Волшебное использование Softmax

Еще одна более ранняя статья«Эффективное внимание: внимание с линейными сложностями»дает более интересный вариант. он замечает$\boldsymbol{QK^{\top}}$середина,$\boldsymbol{Q},\boldsymbol{K}\in \mathbb{R}^{n\times d}$,если"$\boldsymbol{Q}$существует$d$что одно измерение нормализовано, и$\boldsymbol{K}$существует$n$что одно измерение нормализовано", то$\boldsymbol{QK^{\top}}$заключается в автоматическом удовлетворении нормализации, поэтому выбор, который она дает,

в$softmax_1$,$softmax_2$соответственно в первом$(n)$, второе измерение$(d)$Выполните операцию Softmax. То есть в это время каждый из нас дает$\boldsymbol{Q},\boldsymbol{K}$Добавить Softmax, а не закончить$\boldsymbol{QK^{\top}}$Затем добавьте Софтмакс

Фактически можно доказать, что эта форма также является частным случаем формулы (4)​, которая соответствует$\phi(\boldsymbol{q}_i)=softmax(\boldsymbol{q}_i),\varphi(\boldsymbol{k}_j)=e^{\boldsymbol{k}_j}$, читатель может сделать вывод

идея Сушена

Здесь Су Шэнь подал идею. Отправной точкой этой идеи является уже не уравнение (4), а расширение Тейлора исходного определения (2). В расширенном по Тейлору мы имеем

если$\boldsymbol{q}_i^{\top}\boldsymbol{k}_j\geq -1$, то можно гарантировать неотрицательность правой части, так что$\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)=1 + \boldsymbol{q}_i^{\top}\boldsymbol{k}_j$. Читатели, возможно, уже подумали об этом здесь и хотят гарантировать$\boldsymbol{q}_i^{\top}\boldsymbol{k}_j\geq -1$, необходимо только$\boldsymbol{q}_i,\boldsymbol{k}_j$Делать$l_2$Нормализованный. Поэтому окончательный план, предложенный Су Шэнем, таков:

как$\boldsymbol{x}=[x_1,x_2,...,x_n]$,но$\Vert x\Vert=\sqrt{x_1^2+x_2^2+...+x_n^2}$

Это отличается от уравнения (4), но теоретически оно ближе к исходному масштабируемому точечному вниманию.

выполнить

Это в основном для реализации метода, предложенного Су Шеном, но из-за ограниченного уровня автора, в финальном реализованном коде на самом деле есть некоторые проблемы, в основном:

1. По результатам тестов улучшенная скорость расчета не улучшилась
2. Не могу суммировать до 1

Реализация кода предназначена в основном дляPyTorch реализация BERTКод этой статьи, точнее, фактически только измененScaledDotProductAttentionЭто функция, поэтому ниже выложена только эта часть кода

class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def forward(self, Q, K, V, attn_mask):
        Q = F.normalize(Q, dim=3)
        K = F.normalize(K, dim=3)
        M = (torch.ones(Q.shape[0], Q.shape[1], Q.shape[2], K.shape[2]) + torch.matmul(Q, K.transpose(-1, -2))) # scores : [batch_size, n_heads, seq_len, seq_len]
        M_sum = torch.sum(M, dim=3)
        M = M / M_sum.unsqueeze(3).repeat(1, 1, 1, M.shape[3])
        attn = M.masked_fill(attn_mask, 0) # Fills elements of self tensor with value where mask is one.
        context = torch.matmul(attn, V)
        return context

Если у вас есть лучший метод реализации, пожалуйста, дайте мне знать

Reference

• Исследование линейного внимания: должен ли Attention иметь Softmax?