Трансформер для машинного перевода

Заметка о seq2seq оставила немецко-английскую задачу на BLEU 0.152 с энкодером–декодером на LSTM, единственный передаваемый вектор которого терял бо́льшую часть источника. Эта заметка обучает трансформер энкодер–декодер из заметки об архитектуре на тех же данных с теми же токенизаторами и сопоставимым числом параметров, чтобы сравнение выделило именно архитектуру. Затем она разбирает модель на части: что даёт cross-attention, что даёт глубина, что даёт декодирование.

Обозначения. Как и раньше: источник $x_{1:S}$ (немецкий), перевод $y_{1:T}$ (английский), BLEU на тесте считается одной фиксированной реализацией на 1000 тестовых парах, выходы обрезаются по символу конца.

Модель#

nn.Transformer из PyTorch с $d = 128$, 8 головами ($d_k = 16$), $d_{\text{ff}} = 512$, одним блоком энкодера и одним блоком декодера и дропаутом по умолчанию библиотеки. К эмбеддингам токенов с обеих сторон прибавляется выученная позиционная таблица из 32 строк. Для каждого батча строятся три маски:

• маска паддинга источника, чтобы self-attention энкодера и cross-attention декодера игнорировали ключи [PAD] в немецком предложении;
• каузальная маска формы $T \times T$ для self-attention декодера, чтобы позиция перевода $t$ видела только $\le t$;
• маска паддинга перевода для собственного паддинга декодера.

Обучение — teacher forcing со сдвигом на одну позицию и кросс-энтропией без учёта паддинга, ровно как у LSTM; благодаря каузальной маске весь перевод обрабатывается за один проход. Жадное порождение запускает декодер один раз на каждый выходной токен по растущему префиксу и останавливается, когда каждая последовательность в батче выдала [SEP].

Число параметров: 8 385 850 всего против 8 380 474 у LSTM, 497 978 вне общих матриц эмбеддингов и выхода против 492 602. Размеры двух моделей совпадают с точностью до 0.1%.

Результаты#

Модель	Блоки	BLEU на тесте
Энкодер–декодер на LSTM	1 (двунаправленный)	0.152
Трансформер	1 + 1	0.320
Трансформер, декодер видит только первое состояние энкодера	1 + 1	0.190
Трансформер	4 + 4	0.342

Те же данные, те же токенизаторы, тот же бюджет, те же 10 эпох: трансформер вдвое увеличивает BLEU рекуррентной модели. Каждое число — один запуск обучения, оценённый одним батчевым жадным декодером. Второй жадный декодер, лучевой поиск ширины 1 ниже, даёт для того же чекпойнта 4 + 4 значение 0.253 вместо 0.342; две правильные реализации жадного декодирования должны совпадать, так что в одном из двух путей есть невыясненный дефект. Пока он не найден, таблицу стоит читать как один запуск при одном пути оценки: большие разрывы (0.152 против 0.320, 0.190 против 0.320) объясняются им с меньшей вероятностью, чем малые. Примеры предложений показывают, что изменилось:

Источник	Эталон	LSTM	Трансформер (1 + 1)	Трансформер (4 + 4)
Ein Mann mit einem orangefarbenen Hut, der etwas anstarrt.	A man in an orange hat starring at something.	a man in a red shirt is standing on a street.	a man in an orange hat is working on the ground.	a man with an orange hat is talking to something.
Fünf Leute in Winterjacken und mit Helmen stehen im Schnee mit Schneemobilen im Hintergrund.	Five people wearing winter jackets and helmets stand in the snow, with snowmobiles in the background.	three people in a red shirt and a red shirt is sitting on a street.	five people in winter clothing in red and blue snow with snow in the background.	five people in winter hats and wearing a helmet are standing in the snow with snow in the snow.
Leute Reparieren das Dach eines Hauses.	People are fixing the roof of a house.	people are sitting on a street.	people are working on a window of a store.	people are enjoying a window of a train.

Трансформер переносит число (five), цвет (orange), обстановку (snow, winter) и структуру предложения. Он всё ещё путает предметы и действия: крыша становится окном магазина или поезда, починка — работой или удовольствием. Одна гипотеза — эти ошибки идут от содержания, которое корпус из 29 000 пар почти не покрывает, а не от потери информации об источнике по дороге, как у LSTM. Три примера её не проверяют; для этого нужны ошибки, разложенные по частоте слов источника в обучающих данных и по длине предложения.

Если упорядочить четыре результата так, чтобы каждая модель отличалась от предыдущей одним изменением, они раскладывают прирост:

Рисунок

Откуда берётся прирост BLEU

Сколько стоит cross-attention#

Третья строка — абляция, описанная в заметке о внимании: та же модель, у которой cross-attention декодера ограничен первым выходным вектором энкодера, что воссоздаёт узкое место LSTM внутри трансформера. BLEU падает с 0.320 до 0.190. Всё остальное в двух моделях одинаково, поэтому 0.13 из 0.17 прироста относительно LSTM приходится на то, что декодер читает все позиции источника, а оставшиеся 0.04 — на остальную архитектуру (self-attention в обоих стеках, полносвязные слои, динамику обучения).

Сколько стоит глубина#

Четыре блока с каждой стороны поднимают BLEU с 0.320 до 0.342 при 9.77 млн параметров (1.89 млн вне эмбеддингов, почти вчетверо больше, чем у модели с одним блоком). Наблюдаемый выигрыш — 0.022 по одному запуску каждой модели без повторов с другими зёрнами, так что его величина не установлена; правдоподобная причина того, что он мал, — у подписей из 10–15 токенов мало дальней структуры, которую могла бы использовать глубина, но эти запуски её не проверяют. Более глубокая модель обучается без всяких ухищрений, как и должны позволять остаточные связи.

Чему научился cross-attention#

Если извлечь веса cross-attention декодера с одним блоком, получится тензор $[B, T, S]$; каждая строка в сумме даёт 1. Нарисованные для коротких тестовых пар, где строки — английские токены, а столбцы — немецкие, матрицы близки к диагонали со сдвигом на один: строка английского токена $t$ смотрит на немецкий токен, выровненный с английским токеном $t + 1$. Модель предсказывает следующий токен, поэтому ей нужно слово источника, которое переводит следующий выход, — на него она и смотрит. Мягкое выравнивание слов выучено только по парам предложений, без всякой разметки выравнивания.

Лучевой поиск#

Жадное декодирование берёт на каждом шаге самый вероятный токен. Последовательность, которую оно даёт, в общем случае не самая вероятная: чуть менее вероятное первое слово может вести к гораздо более вероятному продолжению. Полный перебор всех $|V|^T$ последовательностей невозможен; лучевой поиск (beam search) хранит на каждом шаге $B$ лучших частичных последовательностей.

На шаге $t$ каждая из $B$ гипотез продолжается своими $B$ самыми вероятными следующими токенами, что даёт $B^2$ кандидатов с оценкой — суммой логарифмов вероятностей $\sum_{s \le t} \log P(y_s \mid y_{<s}, x)$; выживают лучшие $B$. Гипотеза, выдавшая [SEP], завершена и сравнивается с текущей лучшей завершённой по оценке, нормированной на длину,

$$\text{score}(y) = \frac{1}{|y|^{\alpha}} \sum_{s=1}^{|y|} \log P(y_s \mid y_{<s}, x), \qquad \alpha = 0.6,$$

потому что сырые логарифмы вероятностей у более длинных выходов всегда отрицательнее и поощряли бы раннюю остановку. Поиск заканчивается, когда гипотез не осталось или достигнута максимальная длина. Цена — $B$ проходов декодера на выходной шаг вместо одного, и эти проходы нельзя так же просто объединить по гипотезам, как по батчу.

Ширина луча $B$	BLEU на тесте (модель 4 + 4)
1 (жадно)	0.253
3	0.265
5	0.263

Жадное число в этой таблице, 0.253, — не те 0.342, что приведены выше для той же модели; разница 0.089. Они получены разными путями оценки: батчевая оценка порождает с масками паддинга для всего батча, а оценка лучевым поиском кодирует каждое предложение отдельно без паддинга и запускает свой цикл декодирования со своим правилом остановки. Эти различия разрыва не объясняют. Если модель в режиме оценки, а маски, начальный и конечный токены, позиции и очистка вывода одинаковы, жадное декодирование детерминировано, и батч размера 1 и батч с паддингом должны давать одни и те же токены, кроме редких численных равенств. Разрыв такого размера означает, что в одном пути или в обоих есть не найденный дефект. Проверка — сравнить id токенов обоих путей на одних и тех же предложениях, а затем пересчитать все таблицы одним декодером и одной метрикой.

До тех пор эта таблица не показывает, сколько здесь стоит лучевой поиск. Внутри этого одного декодера ширина 3 даёт на 0.012 больше ширины 1, а ширина 5 не выше; это похоже на обычную картину (лучевой поиск немного помогает, быстро насыщается, а более широкие лучи могут даже вредить, потому что нормировка на длину лишь приблизительна; Stahlberg and Byrne, 2019), но декодер, который расходится с другой реализацией своего же случая ширины 1, не доказывает прироста. Числа сравнимы только внутри одной таблицы и одной проверенной реализации — о чём и говорит заметка об оценке.

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

	Энкодер–декодер на LSTM	Трансформер энкодер–декодер
Доступ к источнику при декодировании	Один вектор	Все $S$ выходов энкодера, на каждом шаге, в каждой голове
Последовательных шагов при обучении	$S + T$	0 (маски)
Параметры	8.38 млн	8.39 млн
BLEU на тесте, Multi30k	0.152	0.320 (1 + 1), 0.342 (4 + 4)
Выравнивание	Нет	Читается из cross-attention

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

• Мало данных. 29 000 пар подписей — крошечный объём для перевода; у модели нет знаний за их пределами. Сильные системы в этом семействе задач предобучаются на сотнях миллионов пар предложений или сначала на одноязычном тексте (следующий раздел).
• BLEU. Числа ранжируют эти модели в одном запуске при одном пути оценки, и не более; два жадных декодера расходятся на одном чекпойнте на 0.089, и это расхождение не выяснено.
• Инференс по-прежнему последовательный. Обучение распараллеливается по $T$, порождение — нет, и каждый новый токен заново читает весь префикс, если не кэшировать ключи и значения.
• Смещение экспозиции и разрыв между жадным декодированием и лучевым поиском унаследованы из постановки seq2seq.

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

• При одинаковых параметрах и данных трансформер энкодер–декодер набирает 0.320 BLEU против 0.152 у LSTM; бо́льшая часть прироста (0.13) — доступ декодера ко всем позициям источника.
• Четыре блока с каждой стороны добавляют на этом корпусе 0.022 BLEU в одном запуске — слишком мало данных для устойчивого вывода о приросте; благодаря остаточным связям более глубокая модель обучается так же легко.
• Cross-attention с одним блоком без разметки выучивает диагональное выравнивание со сдвигом на один.
• Лучевой поиск хранит $B$ гипотез с логарифмом вероятности, нормированным на длину; его оценки здесь получены декодером, чей результат при ширине 1 расходится с батчевым жадным декодером (0.253 против 0.342), поэтому прироста от луча они не доказывают.
• Сравнивать BLEU можно только внутри одной реализации и одного пути оценки; если две реализации одного и того же декодирования расходятся, сначала нужно найти дефект, а уже потом читать разницу.

Ноутбук#

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

• Ячейка 58 определяет Seq2SeqTransformer: общий эмбеддинг, выученная позиционная таблица из 32 строк, nn.Transformer(d_model=128, nhead=8, dim_feedforward=512, batch_first=True) с настраиваемым числом блоков, encode с маской паддинга источника, decode с каузальной маской и обеими масками паддинга и жадный generate, который останавливается, когда каждая последовательность выдала [SEP].
• Ячейки 60–64: по одному блоку с каждой стороны, 8 385 850 параметров, 10 эпох; у ячейки обучения нет сохранённого выхода, а текстовая ячейка 64 приводит BLEU на тесте, 0.32. Ячейки 65–70 печатают первые пять переводов.
• Ячейки 72–79 регистрируют forward hook на multihead_attn декодера, предварительно подменив его обёртку так, чтобы она возвращала веса, проверяют форму [32, 31, 32] со строками, дающими в сумме 1, и рисуют шесть матриц выравнивания.
• Ячейки 81–86: вариант без cross-attention; текстовая ячейка 86 приводит его BLEU на тесте, 0.19.
• Ячейки 88–98: по четыре блока с каждой стороны, 9 774 394 параметра, BLEU на тесте 0.342, приведённый в текстовой ячейке 92, и пять переводов.
• Ячейка 100 реализует beam_search для одного закодированного предложения (ширина луча, максимальная длина и $\alpha = 0.6$ — аргументы); список кандидатов обрезается внутри цикла по гипотезам, а не после него, поэтому на каждом шаге соревнуются меньше $B^2$ кандидатов. Ячейки 108–111 оценивают его на тестовой выборке при $B = 1, 3, 5$ и приводят 0.253, 0.265 и 0.263.

Дальше#

Все модели до сих пор обучались с нуля на собственных данных задачи, и лучшая из них, трансформер, сильнее всех выигрывает от большего объёма текста. Следующий раздел курса обучает трансформеры на неразмеченном тексте с целью обучения с самоконтролем и переиспользует их. Заметка о BERT начинает с энкодера: замаскировать токены, предсказать их и получить модель, векторы токенов которой зависят от контекста.

Литература#

• A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, I. Polosukhin. Attention Is All You Need. NeurIPS, 2017. Трансформер энкодер–декодер для перевода; §6.1 (лучевой поиск со штрафом за длину $\alpha = 0.6$).
• Y. Wu et al. Google’s Neural Machine Translation System. arXiv, 2016. §7 (нормировка на длину в лучевом поиске).
• F. Stahlberg, B. Byrne. On NMT Search Errors and Model Errors: Cat Got Your Tongue?. EMNLP, 2019. Почему более широкие лучи перестают помогать.
• D. Elliott, S. Frank, K. Sima’an, L. Specia. Multi30K: Multilingual English-German Image Descriptions. VL, 2016.
• PyTorch. torch.nn.Transformer. Модуль, используемый в ноутбуке; в документации описаны соглашения о масках.