Про расширения RoPE

Первая часть: ссылка



Для начала коротко напомню, что вообще делает RoPE: берётся эмбеддинг и каждая пара фичей вращается на определённый угол. Для разных пар — разный угол. Для первого токена происходит один поворот, для второго — два поворота, и так далее. Вращаем и query, и key.



Кроме того, нас прежде всего интересует расширения контекст без дотюнивания.

Раздобыть данные на понимание длинного контекста и компьют на тюн не так уж просто.





Extending Context Window of Large Language Models via Positional Interpolation

Статья: ссылка



Казалось бы, ну и в чём проблема тогда? Пусть мы обучили модель на 4096 позициях. Тогда давайте просто вращать фичи 4097 раз, 4098 раза и так далее. Это называется экстраполяция 📊, и экспериментально подвтерждено, что она не работает с популярными базовыми моделями. Причём подтверждено ещё много где. При этом со стороны теории это несколько загадочно: RoPE не обучается и кодирует относитетельную позицию, так какого чёрта? То есть ожидается, что после экстраполяции модель будет себя вести так, как будто у неё скользящее окно в 4к токенов. А на самом деле её полностью разносит, логиты внимания для некоторых расстояний >4к улетают в небеса. 🤯

Почему? Об этом чуть ниже.



Пока же сосредоточимся на втором семействе методов — интерполяции. Базовая идея такая: представим, что для токена на 4096 позиции мы делаем в 2 раза меньше оборотов, 2048. И так для каждой позиции, вместо одного оборота делаем половину. После этого мы можем кодировать 8192 токенов вместо 4096. Но есть нюанс: модель-то не видела в обучении полуоборотов, соседние токены становятся практически неразличимы 💀

Поэтому авторы статьи полируют это всё тюном на расширенном контексте, что делает этот подход не особо практичным.





Scaling Laws of RoPE-based Extrapolation

Статья: ссылка



Суть статьи — обоснование провала экстраполяции. А виновата оказалсь база, θ = 10000 😐

Дело в том, что с такой базой не все фичи делают полный оборот за 4к. И в итоге для первых фичей модель видит полный оборот (и видит немонотонность функций), а для других фичей видит куски только до 2π, π или π/2. И как раз из-за наличия неполных кусков экстраполяция не работает как надо.



Авторы дотюнивают модель с разными базами, но в пределах оригинальной длины контекста, 4к. Если сделать базу радикально меньше, например 500, то все фичи совершают по несколько оборотов, и всё ок, экстраполяция будет работать с приемлемым качеством скользящего окна. С увеличением размера контекста становится хуже, но без переломов. Если сделать базу радикально больше, например 1000000, то точка перелома отодвигается на более широкий контекст, но всё ещё существует.



Хоть такой метод и выпадает из нашей изначальной постановки задачи, потому что снова надо тюнить, но тюнить-то надо на маленьком контексте 😊, поэтому проблем со сбором данных тут нет. Работает всё неплохо, особенно с маленькой базой.





NTK-Aware scaling

Пост: ссылка



Меняем интерполяцию с дообучением из первой статьи на масштабирование базы θ без дообучения. Описано в посте на Реддите, хронологически до статьи выше. По идее это всё ещё интерполяция (особенно если задать коэффициент масштабирования с запасом), только мы растягиваем для разных частот по-разному.





Dynamic NTK scaling

Пост: ссылка



Добавляем изменение коэффициента масштабирования в зависимости от текущей длины последовательности. То есть если текущая длины меньше оригинальной — не делаем ничего. А дальше растим коэффициент линейно.





YaRN: Efficient Context Window Extension of Large Language Models

Статья: ссылка



Добавляются две новых штуки: NTK-by-parts и температура в софтмаксе внимания. NTK-by-parts мотивирован фактом выше о том, что какие-то фичи и частоты и так нормально экстраполируются, так что их можно не трогать. Для тех фичей, которые не экстраполируются, мы интерполируем. А температуру меняем, чтобы выравнять распредление скоров внимания (она меньше 1). Да и экспериментально так получается перплексия пониже 🧠



И вот YaRN уже работает как без тюна, так и с тюном.