Архитектура трансформера

Self-attention в том виде, в каком его оставила предыдущая заметка, даёт для каждого токена взвешенное среднее векторов значений. Это операция смешивания по позициям. Веса нелинейно зависят от входа через softmax, но собственные признаки токена ничто нелинейно не преобразует, и нет ни глубины, ни порядка, ни способа порождать. Трансформер (Vaswani et al., 2017) — это self-attention плюс минимум, нужный, чтобы исправить каждое из этого, и в результате получается блок, который можно ставить в стек, обучать параллельно по последовательности и собирать тремя способами в зависимости от задачи.

Обозначения. $x \in \mathbb R^{n \times d}$ — вход блока (по строке на позицию), $d$ — ширина модели, $d_{\text{ff}}$ — ширина полносвязной сети, $L$ — число блоков, $h$ — число голов. $\operatorname{Attn}$ — многоголовый self-attention, а $\operatorname{FFN}$ — полносвязная сеть, применяемая к каждой позиции отдельно (position-wise feed-forward).

Отправная точка: стек слоёв внимания#

Если поставить слои self-attention друг на друга напрямую, сломаются две вещи. Выход каждого слоя линеен по значениям, поэтому $L$ слоёв дают сумму произведений зависящих от содержания матриц усреднения, применённых ко входу, за которыми следуют фиксированные линейные отображения. Веса усреднения зависят от входа нелинейно, но собственные признаки токена никогда не проходят через нелинейную функцию: для входа из одного токена каждая матрица внимания — это скаляр 1, и весь стек — линейное отображение этого токена. Вычислений на уровне отдельного токена нет. К тому же с глубиной такой стек вырождается: без остаточных связей и полносвязных слоёв его выход дважды экспоненциально быстро сходится к матрице ранга 1, и все токены отображаются почти в один и тот же вектор (Dong, Cordonnier and Loukas, 2021). А глубокие стеки любого рода страдают от проблем с градиентом из заметки о RNN: сигнал сверху проходит через $L$ преобразований, прежде чем дойти до низа.

Блок#

При размещении pre-norm (обычный современный выбор) блок записывается так:

$$h' = h + \operatorname{Attn}(\operatorname{Norm}(h)), \qquad h'' = h' + \operatorname{FFN}(\operatorname{Norm}(h')).$$

Рисунок

Блок трансформера: что куда идёт

Шаг за шагом, в порядке номеров на рисунке:

1. Вход. Эмбеддинг каждого токена плюс вектор его позиции даёт $h$, по строке на токен (откуда берётся позиция — тема следующей заметки).
2. Смешать позиции. $h$ нормализуется и проходит многоголовое self-attention — вычисление из предыдущей заметки. Это единственный шаг, на котором одна позиция читает другую.
3. Сложить. Выход внимания прибавляется к $h$: $h' = h + \operatorname{Attn}(\operatorname{Norm}(h))$. Неизменённое $h$ идёт по остаточному пути — синей линии.
4. Преобразовать каждую позицию. $h'$ нормализуется и проходит FFN — одну и ту же двухслойную сеть, применённую к каждой позиции отдельно.
5. Сложить ещё раз. $h'' = h' + \operatorname{FFN}(\operatorname{Norm}(h'))$.
6. Поставить в стопку. У $h''$ та же форма, что у $h$, поэтому оно служит входом следующего блока; через $L$ блоков результат читает голова под задачу.

Полносвязная сеть для каждой позиции#

$$\operatorname{FFN}(x) = \sigma(x W_1 + b_1)\, W_2 + b_2, \qquad W_1 \in \mathbb R^{d \times d_{\text{ff}}},\; W_2 \in \mathbb R^{d_{\text{ff}} \times d},$$

применяется к каждой строке $x$ отдельно с одними и теми же весами. $\sigma$ — ReLU в исходной статье и GELU в BERT и GPT-2; Llama заменяет форму с двумя матрицами на форму с гейтом и тремя матрицами (gate, up, down в списке модулей) и активацией SiLU. Ширину по соглашению берут $d_{\text{ff}} = 4d$: 3072 у BERT-base, 8192 у Llama 3.2 1B шириной 2048 из заметки о GPT. Softmax и LayerNorm тоже нелинейны, но softmax лишь задаёт веса смешивания, а LayerNorm лишь стандартизирует каждый вектор; FFN — обучаемое нелинейное преобразование каждой позиции по отдельности и место, где живёт бо́льшая часть его параметров: $2 d d_{\text{ff}} = 8d^2$ против $4d^2$ у внимания.

Остаточные связи#

$h + f(h)$ вместо $f(h)$. Производная блока $l$ по его входу равна $I + J_l$, где $J_l = \partial f_l/\partial h$. Цепное правило по-прежнему перемножает множители всех $L$ блоков, но произведение раскрывается как

$$\prod_{l=1}^{L} (I + J_l) = I + \sum_{l} J_l + \sum_{l < m} J_m J_l + \dots + J_L \cdots J_1 ,$$

поэтому, помимо слагаемого через все $L$ якобианов, единственного, которое есть у рекурсии, у градиента появляется тождественное слагаемое, не проходящее ни через одну матрицу весов, и короткие слагаемые через один-два блока. Остаточные связи не отменяют цепное правило и не гарантируют устойчивых градиентов; они добавляют короткие пути, и на практике это делает глубокие стеки обучаемыми. Поэтому трансформеры бывают глубокими: 12 блоков у BERT-base, 32 у Llama на 7 млрд параметров, 96 у GPT-3, а заметка о переводе выигрывает 0.02 BLEU в одном запуске, переходя от 1 блока к 4 на 29 000 предложений. Требование одно: вход и выход каждого блока должны иметь одну и ту же ширину $d$, что фиксирует $d$ во всём стеке.

Нормализация#

Нормализация слоя (layer normalization; Ba, Kiros and Hinton, 2016) стандартизует вектор каждой позиции по его $d$ признакам:

$$\operatorname{LN}(x_i) = \gamma \odot \frac{x_i - \mu_i}{\sqrt{\sigma_i^2 + \epsilon}} + \beta, \qquad \mu_i = \frac1d\sum_{j=1}^d x_{ij},\;\; \sigma_i^2 = \frac1d\sum_{j=1}^d (x_{ij} - \mu_i)^2 .$$

В отличие от батч-нормализации (см. заметку курса по генеративным моделям), статистики считаются для каждого примера и каждой позиции, поэтому не зависят от батча и ведут себя одинаково при обучении и инференсе; это важно для последовательностей переменной длины с паддингом. $\gamma, \beta \in \mathbb R^d$ выучиваются. RMSNorm, используемый в Llama, отказывается от вычитания среднего и от $\beta$ и делит только на среднеквадратичное значение; это дешевле и эмпирически не хуже (Zhang and Sennrich, 2019).

Pre-norm и post-norm#

Исходная статья нормализует после остаточной суммы, $h' = \operatorname{LN}(h + \operatorname{Attn}(h))$; это блок post-norm, его использует BERT. Большинство более поздних моделей нормализуют вход каждого подслоя — это блок pre-norm, записанный выше (GPT-2, Llama). Разница в том, куда идёт остаточный путь: в pre-norm это непрерывное тождественное отображение от низа стека до верха, в post-norm он проходит через LayerNorm в каждом блоке. Xiong et al. (2020) показали, что при инициализации блок post-norm даёт большие градиенты параметрам у выхода, а в pre-norm они убывают с глубиной; поэтому моделям post-norm нужны прогрев шага обучения (warm-up) и аккуратная инициализация, а pre-norm обучается устойчиво без прогрева. Post-norm, когда обучается, иногда чуть лучше по итоговой метрике; это обмен устойчивости на небольшую прибавку качества, и на больших масштабах побеждает устойчивость.

Рисунок

Блоки трансформера post-norm и pre-norm

Откуда берётся вход#

Перед первым блоком идентификаторы токенов переводятся в эмбеддинги и к ним добавляется информация о позиции: $x = E[\text{ids}] \cdot \sqrt d + P$, где $P$ кодирует позиции (тема следующей заметки). Множитель $\sqrt d$ согласует масштаб эмбеддинга с масштабом позиционного сигнала. Дропаут применяется к сумме, к весам внимания и после каждого подслоя.

Параметры блока#

На блок, без смещений и векторов нормализации: $4d^2$ во внимании (три проекции плюс выходное отображение) и $8d^2$ в FFN, всего $12d^2$. У небольшого энкодера для AG News ниже $d = 32$ и, что необычно, $d_{\text{ff}} = 32$:

Компонент	Число
Эмбеддинги токенов, $4\,000 \times 32$	128 000
На блок: внимание $4 \times 32^2$	4096
На блок: FFN $2 \times (32 \times 32 + 32)$	2112
На блок: два LayerNorm, $2 \times 64$	128
Четыре блока	25 344
Классификатор $32 \times 4 + 4$	132
Синусоидальная таблица позиций $128 \times 32$ (фиксированная, но зарегистрированная как параметр)	4096
Всего	157 572

Разбивка показывает обычные пропорции: в маленькой модели доминирует таблица эмбеддингов, а блоки крошечные. Для BERT-base ($d = 768$, $L = 12$, словарь с учётом регистра на 28 996 токенов) та же арифметика даёт $12 \times 12 \times 768^2 = 84.9$ млн в блоках, 22.3 млн в эмбеддингах токенов, 0.4 млн в эмбеддингах позиций и 0.6 млн в пулере — всего около 108 млн, именно это число печатает заметка о NER.

Три компоновки#

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

Только энкодер. Блоки с немаскированным self-attention: каждая позиция видит каждую другую в обе стороны. Выход — один контекстный вектор на токен. Используется для задач представления и понимания: классификация через пулинг векторов (модель для AG News ниже усредняет позиции без паддинга; BERT берёт первый токен), разметка токенов, извлечение фрагментов. Главный пример — заметка о BERT.

Только декодер. Блоки с каузальным self-attention: позиция $t$ видит только позиции $\le t$. Выход в позиции $t$ предсказывает токен $t + 1$, так что модель — языковая модель, обучаемая на целых последовательностях за один проход. Господствующая компоновка для порождения; заметка о GPT.

Энкодер–декодер. Стек энкодера над источником и стек декодера, в блоках которого есть третий подслой — cross-attention: запросы приходят от декодера, а ключи и значения — от выхода энкодера; это внимание из заметки о внимании до трансформеров с проекциями. Собственный self-attention декодера каузальный. Естественен для перевода и других преобразований входа в выход; заметка о переводе обучает такую модель.

Рисунок

Энкодер–декодер: декодер читает источник через cross-attention

Три названия описывают поток информации. Токенизация, цель обучения, данные, декодирование и оценка остаются невыбранными, и заметка о семействах моделей сравнивает их по этим осям.

Небольшой энкодер на AG News#

Чтобы поставить число рядом с архитектурой: энкодер из 4 блоков с 4 головами и $d = 32$ над словарём WordPiece из 4000 записей, с усреднением по токенам без паддинга и линейным классификатором, обученный со случайной инициализации 20 эпох с Adam при $10^{-3}$ и размере батча 256, достигает доли правильных ответов на тесте 0.906 на своей лучшей эпохе на AG News. Эта эпоха выбрана по тестовой выборке (ловушка из заметки об оценке), поэтому число немного оптимистично. Базовый уровень TF-IDF из заметки о мешке слов — 0.918, лучшая свёрточная модель — 0.911. Трансформер, обученный с нуля на 120 000 коротких документов, не лучше альтернатив; его преимущество проявляется, когда он предобучен на гораздо большем объёме текста, — это тема следующего раздела курса. Эксперимент показывает другое: архитектура без труда обучается при таком размере, а её карты внимания можно изучать.

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

Рисунок

От RNN к трансформеру: одна идея, шесть доработок

Весь путь от рекуррентной сети к этому блоку, по одному изменению за раз. По сравнению с одним self-attention блок добавляет следующее:

	Только self-attention	Блок трансформера
Вычисления на позиции	Нет (взвешенное среднее)	FFN с $8d^2$ параметрами
Глубина	Произведения матриц усреднения; выход вырождается к рангу 1	$L$ остаточных блоков
Масштаб активаций	Не контролируется	LayerNorm для каждой позиции
Порядок	Нет	Позиционный сигнал добавлен ко входу
Порождение	—	Каузальная маска в декодере
Вторая последовательность	—	Cross-attention

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

• Квадратичное внимание унаследовано; всё остальное в блоке линейно по $n$.
• Фиксированная ширина. Остаточные связи навязывают одно $d$ на весь стек; ёмкость добавляется глубиной или шириной сразу везде.
• Аппетит к данным. При обучении с нуля на малых данных архитектура не превосходит базовые уровни на счётчиках; ей нужно предобучение.
• Позиция — надстройка. Ничто в блоке не знает порядка; качество позиционной схемы ограничивает то, что модель может с ним сделать.
• Число параметров — ещё не архитектура. Две модели с одинаковыми $L, d, h$ могут различаться размещением нормализации, активацией, формой FFN и позиционной схемой, и это меняет поведение при обучении сильнее, чем подсказывают заголовочные числа.

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

• Блок: $h' = h + \operatorname{Attn}(\operatorname{Norm}(h))$, $h'' = h' + \operatorname{FFN}(\operatorname{Norm}(h'))$. Внимание смешивает позиции, FFN преобразует каждую позицию, остаточные связи сохраняют тождественный путь, LayerNorm фиксирует масштаб.
• $12d^2$ параметров на блок ($4d^2$ внимание, $8d^2$ FFN); в маленьких моделях доминирует таблица эмбеддингов.
• Pre-norm: тождественный остаточный путь, устойчивое обучение глубоких стеков. Post-norm: исходный вариант, нужен прогрев.
• Энкодер (двунаправленный), декодер (каузальный), энкодер–декодер (cross-attention): один блок, разные маски.
• С нуля на AG News: 0.906 на лучшей эпохе, ниже TF-IDF; окупиться архитектуре позволяет предобучение.

Ноутбук#

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

• Ячейки 7–27 загружают AG News, обучают токенизатор WordPiece на 4000 записей и строят дополненные батчи формы [256, 128].
• Ячейка 35 реализует синусоидальную таблицу позиций как необучаемый nn.Parameter; ячейки 39–40 — классы внимания и маска паддинга из предыдущей заметки.
• Ячейка 42 определяет TransformerLayer (post-norm, как в исходной статье и BERT: norm1(inputs + attention), затем norm2(outputs + feedforward(outputs)), FFN ширины fc_dim с ReLU и дропаутом) и Transformer (эмбеддинг, умноженный на $\sqrt d$, позиционный кодировщик, четыре слоя, усреднение по маске, линейный классификатор; модель также возвращает собранные веса внимания всех голов и слоёв).
• Ячейки 44–46 создают модель с embed_dim=32, fc_dim=32, num_heads=4, num_layers=4, dropout=0.2 и печатают 157 572 параметра — итог таблицы параметров выше.
• Ячейка 48 обучает 20 эпох с Adam при $10^{-3}$, рисуя кривые функции потерь и доли правильных ответов, и печатает лучшую по эпохам долю правильных ответов на тесте, 0.9064.
• Ячейки 49–62 визуализируют веса внимания.

week04_lstm_attention_transformer.ipynb, ячейка 58, использует nn.Transformer из PyTorch для компоновки энкодер–декодер; её описывает заметка о переводе.

Дальше#

Блок принимает вход как набор векторов с добавленным позиционным сигналом, и пока ничего не сказано о том, что это за сигнал. Следующая заметка показывает, что без него self-attention не отличает «dog bites man» от «man bites dog», выводит синусоидальное кодирование из исходной статьи и сравнивает его с выученными и ротационными альтернативами.

Литература#

• A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, I. Polosukhin. Attention Is All You Need. NeurIPS, 2017. §3.1 (стеки энкодера и декодера), §3.3 (полносвязная сеть), §5.4 (дропаут).
• Y. Dong, J.-B. Cordonnier, A. Loukas. Attention Is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth. ICML, 2021. Rank collapse of attention-only stacks; residuals and MLPs prevent it.
• J. L. Ba, J. R. Kiros, G. E. Hinton. Layer Normalization. arXiv, 2016.
• K. He, X. Zhang, S. Ren, J. Sun. Deep Residual Learning for Image Recognition. CVPR, 2016. Остаточная связь.
• R. Xiong et al. On Layer Normalization in the Transformer Architecture. ICML, 2020. Градиенты pre-norm и post-norm и прогрев.
• B. Zhang, R. Sennrich. Root Mean Square Layer Normalization. NeurIPS, 2019.
• H. Touvron et al. LLaMA: Open and Efficient Foundation Language Models. arXiv, 2023. Pre-norm, RMSNorm, SwiGLU и ротационные позиции в одной модели.
• J. Alammar. The Illustrated Transformer. 2018.