Внимание до трансформеров

Энкодер–декодер из предыдущей заметки сжимает источник в одно состояние фиксированного размера ещё до того, как декодер что-либо напишет, и выходы показывают цену: гладкие предложения не о том. Декодер должен иметь возможность смотреть на источник по ходу письма и смотреть на разные его части для разных слов перевода. Богданау, Чо и Бенджио (2015) добавили именно это и назвали вниманием (attention). Эта заметка строит механизм, вручную проходит один шаг и показывает эксперимент, который измеряет, что он даёт.

Обозначения. Состояния энкодера $h_1, \dots, h_S \in \mathbb R^{H}$ (по одному на токен источника; $2H$, если энкодер двунаправленный). Состояние декодера $s_t \in \mathbb R^{H}$ на шаге $t$. Оценки $e_{t,i}$, веса $\alpha_{t,i}$, вектор контекста $c_t$.

Отправная точка: последнее состояние как единственный канал#

Декодер инициализируется $s_0 = f(h_S)$ и дальше работает сам. Нет операции, которой шаг $t$ декодера мог бы прочитать именно $h_3$. Состояния $h_1, \dots, h_{S-1}$ вычисляются и выбрасываются.

Механизм#

На шаге декодера $t$:

1. Оценить каждое состояние энкодера относительно состояния декодера. Простейшая оценка (attention score) — скалярное произведение, $e_{t,i} = s_t \cdot h_i$; оно велико, когда два вектора смотрят в одну сторону. Есть и другие варианты (ниже), но какой бы ни использовался, на выходе одно число на позицию источника.
2. Нормировать через softmax по позициям источника: $$\alpha_{t,i} = \frac{\exp(e_{t,i})}{\sum_{j=1}^S \exp(e_{t,j})}, \qquad \sum_i \alpha_{t,i} = 1,\ \alpha_{t,i} > 0.$$ Веса образуют распределение по источнику: «сколько от каждого токена источника использовать сейчас».
3. Агрегировать в вектор контекста — взвешенную сумму состояний энкодера: $$c_t = \sum_{i=1}^S \alpha_{t,i}\, h_i \in \mathbb R^{H}.$$
4. Предсказать по состоянию декодера и контексту вместе, например конкатенировав их перед выходным слоем: $P(y_t \mid y_{<t}, x) = \mathrm{softmax}(W_o [s_t; c_t])$, и подать $c_t$ также на следующий шаг декодера.

Каждая операция дифференцируема, поэтому оценивание выучивается обратным распространением ошибки от функции потерь перевода; ничто не говорит модели, какое слово источника соответствует какому слову перевода. Веса пересчитываются на каждом шаге: шаг $t = 1$ может смотреть на первое слово источника, а шаг $t = 5$ — на последнее.

Размерности для одного шага и батча из $B$: состояния энкодера $[B, S, H]$, состояние декодера $[B, H]$, оценки $[B, S]$, веса $[B, S]$, контекст $[B, H]$. По всему переводу длины $T$ веса образуют тензор $[B, T, S]$; именно его показывают тепловые карты в заметке о переводе.

Разобранный пример#

Три состояния энкодера в двух измерениях, $h_1 = (1, 0)$, $h_2 = (0, 1)$, $h_3 = (1, 1)$, и состояние декодера $s = (0.9, 0.2)$:

$i$	$e_i = s \cdot h_i$	$\exp(e_i)$	$\alpha_i$
1	$0.9$	$2.460$	$0.368$
2	$0.2$	$1.221$	$0.183$
3	$1.1$	$3.004$	$0.449$

Знаменатель равен $2.460 + 1.221 + 3.004 = 6.685$. Контекст — $0.368\,(1, 0) + 0.183\,(0, 1) + 0.449\,(1, 1) = (0.817, 0.632)$. Это точка внутри треугольника, натянутого на состояния энкодера, ближе всего к тем, на которые было похоже состояние декодера; ничего не выбрано, всё смешано в пропорции.

Рисунок

Один шаг внимания: что куда идёт

Рисунок проводит четыре шага механизма на этих числах; кнопки под ним показывают их по одному. По сравнению с рисунком энкодера–декодера в энкодере ничего не меняется. Меняется связь энкодера с декодером: вместо одного вектора, переданного один раз, декодер на каждом шаге обращается ко всем состояниям энкодера с весами, которые вычисляет сам.

Функции оценки#

Скалярное произведение предполагает, что $s_t$ и $h_i$ живут в одном пространстве и имеют сравнимый масштаб. Распространены два обобщения (Luong, Pham and Manning, 2015):

Название	$e_{t,i}$	Параметры
скалярное (dot)	$s_t^\top h_i$	нет
общее, билинейное (general)	$s_t^\top W h_i$	$W \in \mathbb R^{H \times H}$
аддитивное (additive; Bahdanau et al., 2015)	$v^\top \tanh(W_1 s_t + W_2 h_i)$	$W_1, W_2, v$

Аддитивная форма была исходной; скалярная — самая дешёвая, и с выученной проекцией обеих сторон именно её использует заметка о self-attention. Косинусное сходство, которое делит на нормы, тоже годится как оценка, но выбрасывает информацию о величине, которой модель может пользоваться; стандартный выбор — ненормированное скалярное произведение, в трансформере масштабированное константой.

Свидетельство: убрать узкое место из трансформера#

Чистое сравнение — не LSTM с вниманием против LSTM без него, потому что эти модели различаются в нескольких отношениях. Более точный эксперимент использует однослойный трансформер из заметки о переводе, декодер которого читает энкодер ровно через описанный выше механизм (там он называется cross-attention). Две версии обучаются одинаково по 10 эпох на Multi30k:

Декодер может смотреть на	BLEU на тесте
все $S$ состояний энкодера	0.320
только первое состояние энкодера (искусственное узкое место)	0.190
энкодер–декодер на LSTM без внимания (предыдущая заметка)	0.152

Ограничение декодера одним вектором энкодера стоит 0.13 BLEU и опускает трансформер почти до уровня LSTM. Всё остальное (self-attention внутри декодера, полносвязные слои, эмбеддинги) не меняется, поэтому разницу можно приписать доступу декодера ко всему набору состояний источника. Проблема — в узком месте, а внимание — его устранение.

В этой задаче выученные веса интерпретируемы. Если нарисовать их матрицей, где строки — токены перевода, а столбцы — токены источника, cross-attention однослойной модели близко к диагонали, причём каждый английский токен смотрит на немецкий токен, следующий за тем, который он переводит: модель предсказывает следующий токен, и самое полезное, что можно прочитать, — слово источника, соответствующее следующему выходу.

Что даёт внимание#

• Нет сводки фиксированного размера. Декодер читает $S$ векторов, какой бы длины ни был источник. Длинные предложения больше не сжимаются сильнее коротких.
• Выравнивание даром. Веса $\alpha_{t,i}$ — мягкое выравнивание слов, выученное только по парам переводов.
• Короткий путь градиента. Функция потерь на шаге перевода $t$ доходит до состояния энкодера $h_i$ через одну взвешенную сумму, а не через $S + t$ рекуррентных шагов, что облегчает проблему затухания из заметки о RNN.

Что изменилось#

	Энкодер–декодер	С вниманием
Информация об источнике на шаге	$s_0$, фиксированная	$c_t$, пересчитывается
Используемые состояния энкодера	только $h_S$	все $h_1, \dots, h_S$
Дополнительные вычисления на шаге	—	$S$ оценок, один softmax, одна взвешенная сумма
Дополнительные параметры	—	нет (dot) до $O(H^2)$ (general, additive)
BLEU, Multi30k, декодер однослойного трансформера	0.190 (с узким местом)	0.320

Ограничения#

• Энкодер по-прежнему рекуррентный. Состояния $h_i$ вычисляются последовательно, поэтому энкодер не распараллеливается вдоль источника, и каждое $h_i$ по-прежнему несёт источник только через цепочку рекурсий.
• Внимание только между декодером и энкодером. Внутри источника токен $i$ узнаёт о токене $j$ только через рекурсию; внутри перевода — так же. Нет причин, по которым та же оценка не могла бы связывать позиции одной последовательности между собой.
• Цена шага — $O(S)$. Для предложений нормально; для документов уже заметно, а когда каждая позиция смотрит на каждую, цена становится квадратичной.
• Мягкие веса — не объяснения. Почти диагональное выравнивание на наборе подписей к изображениям не обобщается до утверждения, что веса внимания показывают, что модель «использует» (Jain and Wallace, 2019).

Что нужно запомнить#

• Внимание на шаге $t$: оценки $e_{t,i}$ между состоянием декодера и каждым состоянием энкодера, softmax по $i$, контекст $c_t = \sum_i \alpha_{t,i} h_i$, предсказание по $[s_t; c_t]$.
• Веса — распределение по позициям источника, пересчитываемое для каждого выходного токена и выучиваемое по функции потерь задачи.
• Оценка скалярным произведением самая дешёвая; аддитивный и билинейный варианты добавляют параметры.
• Если убрать у декодера доступ ко всем состояниям энкодера, в остальном идентичная модель теряет на Multi30k 0.13 BLEU.
• Внимание убирает узкое место, но не рекурсию.

Ноутбук#

week04_lstm_attention_transformer.ipynb. Номера ячеек считают все ячейки, включая текстовые, начиная с 1. Приведённые числа взяты из прогона, сохранённого в курсе; у ячеек, изменённых с тех пор под текущие версии библиотек, сохранённого выхода нет.

• Ячейки 58–64 определяют и обучают однослойный Seq2SeqTransformer (скрытый размер 128, 8 голов); у ячейки обучения нет сохранённого выхода, а текстовая ячейка 64 приводит BLEU на тесте, 0.32. Модель описывает заметка о переводе.
• Ячейки 72–79 извлекают веса cross-attention декодера с помощью forward hook (тензор [32, 31, 32] весов батч × перевод × источник, каждая строка в сумме даёт 1) и рисуют их для шести коротких тестовых пар; описанный выше почти диагональный шаблон со сдвигом на один токен виден на сохранённом рисунке.
• Ячейки 81–86 определяют Seq2SeqTransformerNoCA, чей decode передаёт в качестве памяти энкодера только encoded_input[:, :1] (состояние первого токена источника), обучают её так же; текстовая ячейка 86 приводит её BLEU, 0.19.

Дальше#

Внимание сравнивало состояние декодера с состояниями энкодера. Ничто в механизме не требует, чтобы стороны были из разных сетей: последовательность может смотреть на саму себя, когда каждая позиция оценивает каждую другую и агрегирует их значения. Так рекурсия исчезает совсем, потому что новое представление каждой позиции — взвешенная сумма, вычисленная одним матричным произведением параллельно по всей последовательности. Следующая заметка записывает self-attention в матричной форме, с запросами, ключами и значениями, и на нём строится трансформер.

Литература#

• D. Bahdanau, K. Cho, Y. Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. ICLR, 2015. Аддитивное внимание.
• M.-T. Luong, H. Pham, C. D. Manning. Effective Approaches to Attention-based Neural Machine Translation. EMNLP, 2015. Оценки dot, general и concat; глобальное и локальное внимание.
• S. Jain, B. C. Wallace. Attention is not Explanation. NAACL, 2019.
• D. Jurafsky, J. H. Martin. Speech and Language Processing, 3rd ed. draft. Глава 13 (машинный перевод; энкодер–декодер с вниманием).