О методах позиционного кодирования в Transformer12.12.2023 16:30

Традиционный дисклеймер

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

Введение

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

В общем случае Transformer, в отличие от, например, RNN, обрабатывает входные векторы одновременно и без дополнительной информации о позиции каждого токена будет рассматривать последовательность как «мешок слов». Позиции могут кодироваться различными способами: с использованием синусоидальных векторов и без них, с добавлением векторов или модернизацией подсчёта внимания, с экстраполяцией или интерполяцией обученного контекста. Важными требованиями являются уникальность представления для каждой позиции, независимость расстояний между парами токенов от длины входа, детерминированность представлений и адаптация метода к расширению контекста модели.

Разберём основные популярные подходы, описанные в научных работах и используемые в известных LLM. Изображения либо нарисованы для статьи, либо взяты из первоисточников, туда же рекомендуется идти за подробностями реализаций и экспериментов. Названия подходов взяты либо из самих работ, либо введены в статье и являются сквозными для удобства навигации.

Основные обозначения и формулы

$\qquad e_{ij} = \cfrac{\big(x_i W^Q\big)\big(x_j W^K\big)^T}{\sqrt{d_z}}$ $\qquad \alpha_{ij} = \mathrm{softmax}(e)_{ij} = \cfrac{\exp(e_{ij})}{\sum_{k=1}^n \exp(e_{ik})}$ $z_{i} = \sum_{j=1}^n \alpha_{ij}\big(x_{ij}W^V\big)$

Обзор работ

Sinusoidal и Trainable Absolute, 2017

Первый, наиболее простой, но вполне рабочий способ, был предложен авторами самого Transformer. Позиционная информация кодируется векторами, по одному на каждую абсолютную позицию в обрабатываемом контексте. Вектор позиции прибавляется к смысловому вектору токена (эмбеддингу) на входе модели, после чего результат подаётся в основную часть нейросети.

Сформировать вектор $\hat p_t \in \mathbb{R}^d_x$ для позиции можно разными способами. Один из вариантов — простой учить такие векторы вместе с остальными частями сети. Это увеличивает число параметров и, как следствие, для небольших моделей может вырастить объём вычислений. Кроме того, при таком подходе модель не может обрабатывать последовательности с длиной, большей её контекста на обучении.

Более популярный и аналогичный по качеству работы метод — использование синусоидальных (тригонометрических) векторов. Элементы $\hat p^i_t$ вектора определяются формулой:

$\hat p^i_t = \begin{cases} \sin(w_k t), \quad i = 2k \\ \cos(w_k t), \quad i = 2k + 1 \\ \end{cases}\hspace{-10pt}, \qquad w_k = \cfrac{1}{10000^{2k / d_x}}$

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

Relative, 2018

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

Эта работа была одной из первых в области, в ней предлагается для каждого расстояния i - j обучать два вектора $a_{ij}^K, a_{ij}^V \in \mathbb{R}^{d_z}$ и использовать их в обновлённых формулах подсчёта self-attention:

$e_{ij} = \cfrac{\big(x_i W^Q\big)\big(x_j W^K \color{blue}{\,+\,a_{ij}^K}\big)^T}{\sqrt{d_z}}, \qquad z_{i} = \sum_{j=1}^n \alpha_{ij}\big(x_{ij}W^V\color{blue}{\,+\,a_{ij}^V} \big)$

Информация о позиции вкладывается в векторы ключей и значений. Вычисления можно ускорить, раскрыв скобки перед подсчётом, бученные векторы общие для всех голов внимания и слоёв.

Transformer-XL, 2019

Вход модели делится на сегменты, которые обрабатываются последовательно, при обработке -го сегмента используются выходы для i+1 -го. Абсолютное позиционное кодирование работать не будет: у обоих сегментов оно одинаковое. Предлагается относительная схема, как и в Relative, позиционная информация переходит в self-attention. Представим x_i в виде суммы смысловой и позиционной частей $x_i = x_i^{\textbf{e}mb} + x_i^{\mathbf{p}os}$ и распишем формулу подсчёта логитов $e_{ij}$ :

$e_{ij} \propto \big((x^e_i \,+\, x^p_i) W^Q\big)\big((x^e_j \,+\, x^p_j) W^K \big)^T = \\ \qquad\qquad\qquad = x^e_i W^Q \big(W^K\big)^T \big(x^e_j\big)^T \,+\, x^e_i W^Q \big(W^K\big)^T \big(x^p_j\big)^T \,+ \\ \qquad\qquad\qquad +\, x^p_i W^Q \big(W^K\big)^T \big(x^e_j\big)^T \,+\, x^p_i W^Q \big(W^K\big)^T \big(x^p_j\big)^T$

Репараметризуем полученное разложение для относительного кодирования, избавившись от зависимости от компоненты x^p_i :

$e_{ij} \propto \underset{(1)}{\underbrace{x^e_i W^Q \big(\color{blue}{W^{K, E}}\big)^T \big(x^e_j\big)^T}} \,+\, \underset{(2)}{\underbrace{x^e_i W^Q \big(\color{blue}{W^{K, R}}\big)^T \big(\color{blue}{ R_{i-j}}\big)^T}} \,+ \\ +\, \underset{(3)}{\underbrace{{\color{purple}u} \big(\color{blue}{W^{K, E}}\big)^T \big(x^e_j\big)^T}} \,+\, \underset{(4)}{\underbrace{{\color{purple}v} \big(\color{blue}{W^{K, R}}\big)^T \big(\color{blue}{ R_{i-j}}\big)^T}}$

Абсолютные позиционные эмбеддинги заменяются на фиксированную синусоидальную матрицу относительных ${\color{blue}R}$ . Векторы запросов для позиционной предлагается брать одинаковыми и моделировать обучаемыми векторами ${\color{purple}u}$ и ${\color{purple}v}$ . Выделяются две весовые матрицы для получения векторов ключей эмбеддингов и позиций соответственно.

Интуитивный смысл каждого из слагаемых можно описать так:

содержательная смысловая связь
связь смысла и позиционной информации
глобальный смысловой сдвиг
глобальный позиционный сдвиг

В Relative есть только (1) и (2), сдвиги отбрасываются, хотя используются дополнительная модификация векторов значений. Кроме того, в Relative произведение $\big(\color{blue}{W^{K, R}}\big)^T \big(\color{blue}{ R_{i-j}}\big)^T$ заменяется одной обучаемой матрицей, что означает отказ от синусоидальных векторов и, как следствие, ухудшение обобщения на длинный контекст.

T5, 2020

В этой популярной модели позиционная информация кодируется скаляром, который прибавляется к $e_{ij}$ перед softmax. Скаляры соответствуют различным расстояниям-отступам (i-j) . Всего скаляров 32, они в логарифмической шкале покрывают 128 отступов, т.е. близкие к 0 соседние отступы кодируются разными скалярами, а далёкие могут кодироваться одним:

Снизу номер скаляра, внутри — соответствующие ему расстояния между токенами.

Всем расстояниям, большим 128, соответствует один и тот же скаляр, он способствует обобщению модели на длинный контекст. Позиционные скаляры обучаемые и настраиваются вместе с моделью. Каждый следующий слой расширяет окно улавливаемой позиционной информации (идейно схоже с CNN или окном внимания в Longformer). Обученные скаляры свои для каждой головы self-attention, но общие для всех слоёв модели.

DeBERTa, 2021

Как и в Transformer-XL, формула подсчёта логитов в self-attention раскладывается, и предлагается использовать первые 3 слагаемых:

$e_{ij} \propto \underset{(1)}{\underbrace{x^e_i W^Q \big(W^K\big)^T \big(x^e_j\big)^T}} \,+\, \underset{(2)}{\underbrace{x^e_i W^Q \big(W^K\big)^T \big(x^p_j\big)^T}} \,+ \\ +\, \underset{(3)}{\underbrace{x^p_i W^Q \big(W^K\big)^T \big(x^e_j\big)^T}} \,+\, \underset{(4)}{\underbrace{\color{red}{{x^p_i W^Q \big(W^K\big)^T \big(x^p_j\big)^T}}}}$

По мнению авторов, слагаемое (3) тоже является важным важным для более полного моделирования информации о сдвиге, в то время как (4) — position-to-position term — для относительного кодирования несёт мало дополнительной информации.

Предлагаемая репараметризация довольно похоже на Transformer-XL:

$e_{ij} \propto \underset{(1)}{\underbrace{ x^e_i W^Q \big(W^K\big)^T \big(x^e_j\big)^T }} \,+\, \underset{(2)}{\underbrace{ x^e_i W^Q \big(\color{blue}{W^{K, R}}\big)^T \big(\color{blue}{ R_{i-j}}\big)^T }} \,+ \, \underset{(3)}{\underbrace{ \color{blue}{R_{i-j}} \color{blue}{W^{Q, R}} \big(W^K\big)^T \big(x^e_j\big)^T }}$

Новые весовые матрицы свои у каждой пары голова-слой, матрица расстояний $R_{i-j}$ обучаемая, общая для всех слоёв и своя у каждой головы. Для всех отступов, не попадающих в размер $R_{i-j}$ , берутся векторы ближайшего с соответствующего конца отступа.

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

RoPE, 2021

Один из самых популярных методов, который не только регулярно используется в современных моделях (LLaMA, Qwen, Mistral), но активно модернизируются, о чем речь будет идти ниже.

Рассмотрим метод шаг за шагом. Если исключить позиционные эмбеддинги, то в обобщенном варианте формула для логитов выглядит так

$e_{ij} \propto \langle f_q(x_i, i), f_k(x_j, j) \rangle,$

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

$\langle f_q(x_i, i), f_k(x_j, j)\rangle = g(x_i, x_j, i - j)$

Можно доказать, что для случая d_z=2 подойдёт преобразование

$f_q(x_i, i) = \begin{pmatrix} \cos i\theta & -\sin i\theta \\ \sin i\theta & \cos i\theta \end{pmatrix} \begin{pmatrix} W^Q_{11} & W^Q_{12} \\ W^Q_{21} & W^Q_{22} \end{pmatrix} \begin{pmatrix} x_i^1 \\ x_i^2 \end{pmatrix},$

где $\theta \in \mathbb{R}, \theta \ne 0$ — заданная константа. Формула для f_k аналогична. Подходящая функция при этом определяется так:

$g(x_i, x_j, i - j) = \mathrm{Re}[\big( W^Q x_i \big)\big( W^K x_j \big)^* e^{\textbf{i}(i - j)\theta}],$

где $\mathrm{Re[X]}$ — вещественная часть $X \in \mathbb{C}$ , A^* — сопряжённая матрица к , а $\textbf{i}$ — мнимая единица. Применение RoPE заключается в повороте вектора запроса/ключа на угол, зависящий от индекса его позиции. Поворот обоих векторов на один угол, т.е. смещение позиций без изменения расстояния, сохранит значение скалярного произведения.

Полученное преобразование обобщается на любую чётную d_z . Для этого d_z делится на d_z / 2 двумерных подпространств, к каждому из которых применяется свой матрица поворота. Итоговое выражение для f_q ( f_k аналогично):

$f_q(x_i, i) = R^{d_z}_{\Theta, i} \, W^Q \, x_i,$

где

$R^{d_x}_{\Theta, i} = \begin{pmatrix} \cos i\theta_1 & -\sin i\theta_1 & 0 & 0 & \cdots & 0 & 0 \\ \sin i\theta_1 & \cos i\theta_1 & 0 & 0 & \cdots & 0 & 0 \\ 0 & 0 & \cos i\theta_2 & -\sin i\theta_2 & \cdots & 0 & 0 \\ 0 & 0 & \sin i\theta_2 & \cos i\theta_2 & \cdots & 0 & 0 \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & \cdots & 0 & 0 & \cos i\theta_{d_z / 2} & -\sin i\theta_{d_z / 2}\\ 0 & 0 & \cdots & 0 & 0 & \sin i\theta_{d_z / 2} & \cos i\theta_{d_z / 2}\\ \end{pmatrix},$ $\Theta = \{\theta_k = 10000^{-2(k-1)/d_z}, k \in [1, \dots, d_z / 2] \}$ Визуализация применения RoPE.

Визуализация применения RoPE.

В отличие от Absolute применяется не отдельным координатам, а к парам, и использует умножение на $\sin/\cos$ вместо суммы (мультипликация вместо аддитивности). В разных экспериментах RoPE показывает себя лучше, чем Absolute и Relative:

Сравнение моделей с Absolute и RoPE по значению функционала качества на задачах MLM и языкового моделирования.

Теоретически RoPE должен помогать модели обобщаться на более длинный контекст, чем при обучении, но по факту это не работает, поэтому предпринимаются попытки улучшить его.

ALiBi, 2022

Этот метод тоже является достаточно популярным и развивает подход, схожий с T5: вместо добавления позиционных эмбеддингов или репараметризации подсчёта внимания к $e_{ij}$ добавляется скаляр. Но в отличие от T5 скаляры не обучаются и представляют собой произведение m(i - j) , где — заданное на старте число, своё для каждой головы и общее для слоёв. Например, для 8 голов это могут быть $\frac{1}{2^1}, \frac{1}{2^2}, \dots, \frac{1}{2^8}$ .

В сравнении Absolute vs. RoPE vs T5 только T5 показал хорошую способность к обобщению на последовательностях с большей длиной, чем на обучении (до 600 токенов сверх исходных 512), но его преимущество перекрывается возрастающими вычислительными затратами, проще обучить модель с Absolute с большим контекстом. AliBi оказывается достаточно эффективным с т.з. скорости и памяти:

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

xPos, 2023

Один из способ доработки RoPE для лучшей адаптации к длинному контексту. Вводится понятие «ожидаемого значения внимания» для пары токенов на заданном расстоянии. Показывается, что в обычном RoPE при существенном росте расстояния эта величина начинает осциллировать, это портит качество модели:

Проблема объясняется тем, что значения косинуса не являются монотонными при угле поворота, большем $\pi$ . Решение: добавить задаваемые априорно дополнительные множители для каждой пары компонентов векторов и ( и $T^{-1}$ в листинге алгоритма подсчёта внимания из статьи):

Несмотря на отличия в нотации, несложно узнать в описанном алгоритме RoPE с дополнительными множителями в формулах запросов и ключей. M — маска внимания.

Введённые множители масштабируют компоненты векторов и стабилизируют график ожидаемого значения внимания, улучшая поведение модели на длинном контексте.

NoPE, 2023

В введении было сказано, что отсутствие позиционной информации превращает Transformer в модель «мешка слов», но это точно верно только для моделей на основе кодировщика. Декодировщик, обучаемый авторегрессионно, теоретически способен работать без явной позиционной информации по аналогии с рекурентными сетями.

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

Для проверки гипотезы предлагается маленький эксперимент с моделью размера 100M , контекстом на обучении 20 и на тесте до 40. Модели с разным позиционным кодированием обучаются на 10 задачах, разделенных на 3 группы (задачи вида «сложить 2 числа» или «переставить слова»), можно видеть, что NoPE показывает себя неплохо:

Авторы задаются вопросом о схожести соответствующих выходов слоёв моделей для разных подходов. Близость между моделями и на слое $\ell$ можно определить как

$D^\ell(A, B) = \underset{(P, Q) \in A_\ell \times B_\ell}{\min} D(P,Q), \quad D(P,Q) = \cfrac{1}{n}\sum_{i=1}^nD_{JSD}(P_i \Vert Q_i),$

где $D_{JSD}$ — дивергенция Йенсена-Шеннона по выходам двух голов. Оказалось, что модель, обучаемая без явной позиционной информации, более всего похожа на модель с кодированием T5:

Второй эксперимент ближе к реальной жизни: модель 1.3B , контекст 1024 (на тесте до 2560), обучение авторегрессионное на данных из StarCoder. На графиках показана перплексия на 200 последних токенах:

Сверху — длина исходных текстов.

NoPE всё равно ломается при росте длины контекста, хотя показывает себя ощутимо лучше, чем RoPE.

Positional Interpolation RoPE, 2023

Эта работа была опубликована почти одновременно с SuperHOT RoPE, в обеих описываются схожие идеи, существенно изменяющие идею расширения контекста обученной модели. Описанные выше подходы делают экстраполяцию: обучение на одном диапазоне, работа на другом. Вместо этого можно делать интерполяцию: обучить модель с RoPE и вложить увеличенный контекст в тот же диапазон с минимальным дообучением.

Вместо экстраполяции можно сжать уже обученный диапазон.

Вспомним функцию для получения $e_{ij}$ из RoPE и определим вместо неё новую функцию :

$g'(x_i, x_j, i - j) = g\bigg(x_i, x_j, \frac{(i - j)\color{blue}{L}}{\color{blue}{L'}}\bigg),$

где и — исходная и увеличенная длины контекста. Т.е. если модель выучила позиции 1, 2, 3 и т.д., сжатие в пределах диапазона этих значений приведёт к переходу на кодирование позиций 0.5, 1, 1.5, 2 и т.д.

Показано, что для адаптации к новому контексту достаточно дообучения на $\sim 10^4-10^5$ последовательностях, а замеры перплексии и качества на части бенчмарков LLaMA и суммаризации показывают преимущество 1К шагов дообучения с PI над 10К шагами обычного дообучения (FT).

NTK-Aware Scaled RoPE, 2023

В PI RoPE предлагается по сути делать линейную интерполяцию, что не является оптимальным вариантом. Согласно обзорному анализу из работы YaRN на основе теории Neural Tangent Kernel (NTK), для DL-моделей задача выучить многомерный комплексный вектор для кодирования одномерной по своей сути позиционной информации является проблемной. RoPE имеет сходство со специальным одномерным видом Fourier Features из NTK, и равномерное растяжение векторов RoPE (как в PI) приводит к потере высокочастотных деталей, необходимых для различения очень похожих и близких в тексте токенов. Т.е. минимальный поворот, отличающий позиции, оказывается слишком маленьким. Это может быть причиной некоторого падения качества PI RoPE на не-длинных сэмплах после дообучения (чего быть не должно).

В качестве одного из возможных альтернатив предлагается вместо изменения масштаба менять основание (которое 10000), и, как следствие, «скорость вращения» векторов:

$b_{new} = b\, \cdot\color{blue}{ \bigg(\frac{L'}{L}\bigg)^{\frac{d_z}{d_z - 2}}}$

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

Схожий подход использовался при обучении моделей LLaMA Long.

NTK-by-parts RoPE, 2023

Довольно быстро был предложен метод, нивелирующий проблемы двух описанных выше. Введем определение длины волны — числа токенов, необходимого для полного прохода RoPE круга $2\pi$ в измерении :

$\lambda_d = \cfrac{2\pi}{\theta_d}$

PI RoPE и NTK-Aware RoPE не учитывают длину волны и считают все измерения RoPE одинаково важными для модели. Наблюдения же показывают, что у части измерений $\lambda_d > L» src=«https://habrastorage.org/getpro/habr/upload_files/534/00a/1ca/53400a1ca4359d5b5089e0af3829c93c.svg» />, у части наоборот, и дисбаланс может быть сильным. С каждым измерением можно работать по-своему в зависимости от его длины волны. В методе NTK-by-parts RoPE предлагается: </p> <ul><li><p>если <img alt=$ — не интерполировать

если $\lambda_d \ge L$ — интерполировать без экстраполяции (как в PI RoPE)

для прочих делать обычную NTK-Aware интерполяцию

Полученный подход работает лучше PI RoPE и NTK-Aware и с дообучением, и без него.

Dynamic Scaled RoPE, 2023

Ещё одна идея в копилку интерполяционных методов, суть которой в том, чтобы использовать не фиксированный фактор масштаба, а динамический, величина которого зависит от длины $\ell$ текущей обрабатываемой последовательности. Т.е. вместо L' / L используется $\ell' / L, \quad\ell'=\max(1, \ell / L)$ . Показано, что этот подход повышает устойчивость модели к изменению длины контекста.

YaRN, 2023

Работа прежде всего интересна структурированным обзором основных интерполяционных методов (описанных выше), но также авторы предлагают модификацию NTK-by-parts RoPE. Утверждается, что добавление температуры в знаменатель формулы для $e_{ij}$ хорошо влияет на перплексию при расширении контекста. Предлагаемый YaRN = NTK-by-parts RoPE + температура при подсчёте логитов.

В наборе экспериментов подход показывает хорошие результаты (как по перплексии, так и на бенчмарках), кроме того, он (как и другие интерполяционные методы) несложно реализуется и совместим с оптимизациями подсчёта внимания, например, с Flash Attention.

Заключение

Тема кодирования позиционной информации в тренде, работ выходит много и в ближайшем будущем поток не планирует иссякать. Некоторые идеи реализуются в общепринятых фреймворках ещё до оформления в виде научных публикаций, как это случилось с модификациями интерполяционных кодирований (например, класс LlamaDynamicNTKScalingRotaryEmbedding в библиотеке transformers).

Можно точно сказать, что однозначного стандарта в этой области ещё нет, хотя чаще всего современные модели учатся на основе ротационных эмбеддингов RoPE и их модификаций, за что можно поблагодарить успех LLaMA. Интерполяционные методы RoPE выглядят очень изящно и перспективно, интересно будет продолжать наблюдать за их развитием. Но пока основной рецепт при обучении своей модели с нуля — попробовать 2–3 популярных варианта на небольшом числе параметров и принимать решение исходя из полученных результатов.

Спасибо за внимание и успехов!