Всем привет!

Мы в команде фундаментальных исследований MTS AI занимаемся исследованиями в области обработки естественного языка и компьютерного зрения, а также строим свои фундаментальные языковые модели. Недавно у нас получилось достичь уровня gpt-4 на собственном ограниченном датасете большого контекста. Расскажем, как нам это удалось.

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

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

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

Помимо стандартных инструкционных моделей, наша команда в MTS AI разрабатывает LLM с длинным окном контекста – наши последние такие модели мы назвали Cotype Plus 16k и Cotype Plus 32k. Для обучения мы собрали датасет, состоящий из длинных документов и инструкций к ним, на которых впоследствии обучили наши базовые модели. Основная трудность в обучении таких моделей заключается в вычислительных затратах - механизм внимания (self-attention) имеет квадратичную сложность, и, например, обучение модели на контексте 8192 требует в 16 раз больше ресурсов чем для контекста 2048. 

Существует ряд способов адаптации моделей для работы с длинным контекстом. Как правило это:

1. модификации в архитектуре трансформера;

2. различные способы построения и интерполяции позиционных эмбеддингов; 

3. модификации окна контекста;

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

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

Эффективные механизмы внимания

Разберем, как решить основную проблему механизма внимания при работе с большим контекстом – квадратичную сложность при его вычислении. Больших затрат на вычисления можно избежать, используя модификации механизма внимания – это могут быть его локальные или разреженные вариации. Например, COLT5 [1] имеет в архитектуре две ветви: с обычным механизмом внимания и его легковесными вариациями, LongNet [2] использует dilated attention, а LongT5 [3] предлагает временное глобальное внимание (TGlobal), применяя механизм внимания к локальным и глобальным токенам (усреднение токенов в окне).

Такие методы, как правило, требуют дополнительного обучения модели.

Для обучения наших моделей мы использовали LongLoRA. Данный фреймворк представлен в статье “Efficient fine-tuning of long-context large language models” [4], где авторы предлагают использовать shifted sparse attention - модификацию стандартного self-attention. Этот метод предполагает что контекст делится на несколько групп, для каждой из которых внимание считается отдельно. Каждый attention head распадается на две части, и в одной из них токены смещаются, что позволяет информации передаваться из одной группы в другую, при этом значительно сократить количество вычислений. При этом во время инференса используется dense attention, что позволяет модели достигать такого же качества генераций, как и после full-attention файнтюнинга.

Позиционная интерполяция

Другой ряд методов – адаптирующих позиционные эмбеддинги – позволяет запускать на длинных задачах модели, обученные на коротких инструкциях.Такие RoPE Scaling подходы можно применить для увеличения длины контекста модели либо во время инференса без необходимости дополнительного дообучения (либо с коротким SFT этапом).

Сам RoPE (Rotary Position Embedding), впервые описанный в “Enhanced transformer with rotary position embedding” [5], очень прост: attention score между двумя токенами, показывающий насколько много “внимания” передается от одного токена к другому,  должен зависеть только от относительного расстояния между ними.  Таким образом, функция скора зависит только от векторных представлений двух токенов и относительного расстояния между ними. 

Благодаря относительной природе таких позиционных эмбеддингов, их можно скалировать, выходя за пределы изначального максимального контекста. В [6] предлагают модифицировать функцию, линейно «втиснув» новую последовательность в исходное контекстное окно. 

В NTK-aware интерполяции предлагают вместо того чтобы масштабировать частоты каждого измерения RoPE на коэффициент 1/s, распределять их интерполяцию по нескольким измерениям, масштабируя высокие частоты меньше, а низкие частоты больше. Есть и другие, более сложные, способы скейлинга – YaRN [7], динамический скейлинг; однако для наших моделей мы решили использовать простой линейный скейлинг.

Наша модель Cotype Pro 16k изначально была обучена на контекст 8k, однако поскольку она имеет в архитектуре RoPE, мы смогли из коробки увеличить контекст в два раза. 

Оценка long-context моделей

Для оценки наших моделей мы составили два бенчмарка (на контексты 16к и 32к) по 50 вопросов, в которых собрали инструкции с длинным контекстом, из которого модель впоследствии должна достать необходимую информацию. Пример вопроса с ответами нашей модели и бейзлайнов приведен ниже.

Под какую бы задачу ни обучалась модель, часто встает вопрос - как именно ее оценивать? В зависимости от типа задачи могут использоваться или автоматические метрики, или мануальная оценка экспертами. В бенчмарке LongBench, например, для оценки качества англоязычных моделей на длинном контексте используют статистические метрики ROUGE-L и F1. Однако использование таких метрик подходит только для определенного класса вопросов, поскольку они не учитывают синонимичность ответов и их смысловую подоплёку в целом. Чтобы с этим бороться, для оценки больших языковых моделей прибегают к помощи человека – например, при сравнении нескольких языковых моделей часто используют side-by-side метод, когда эксперт сравнивает два ответа и говорит, чей лучше. Такой метод оценки считается очень качественным, но дорогим и долгим. 

Схематичное представление работы модели для получения оценки e5m-score.Для оценки наших моделей мы решили использовать эмбеддинг-модель e5-mistral. Для получения численного значения, показывающего соответствие сгенерированного ответа “ground-truth” ответу, мы прогоняем их через e5-mistral и получаем два эмбеддинга. При перемножении двух векторных представлений мы получаем оценку их “похожести” - значения от 0 до 1, которые затем можно усреднить, пройдясь по всем вопросам из бенчмарка. Такую метрику будем называть e5m-score.

Для выбора метода оценки моделей с большим окном контекста мы провели небольшой эксперимент. Мы сравнили метрики для вычисления “похожести” пар предложений, посчитав корреляцию с распределением оценок от команды экспертов. 

Команда экспертов давала каждому ответу модели тернарную оценку - 0, если ответ плохой и не похож на ground truth; 0.5 - ответ неплохой, но в чем то не соответствует ground truth; 1 - ответ хороший и соответствует ground truth. Затем оценки для каждого вопроса усреднялись по всем экспертам.

Среди представленных методов есть:

1. описанный ранее e5m-score

2. sbert

3. bertscore

4. метрика gpt4 - просим численно (от 0 до 1) оценить насколько похож ответ модели на ground-truth

5. MIS - mutual implication score

Метод оценки моделью e5 выигрывает остальные методы, показывая корреляцию 0.8952 по Пирсону с экспертной оценкой (EV).

Ниже приведена диаграмму где сравниваются различные модели  по нашему бенчмарку 16k. Чем больше значение – тем лучше результат.

Используя описанный ранее подход, мы сперва обучили модели Cotype Plus 16k  и Cotype Plus 32k с помощью LongLoRA. Данные модели показали себя хорошо на наших бенчмарках, держась на уровне с gpt-3.5-turbo, однако не догоняя gpt-4 (у модели заявлен контекст 64k токенов). Однако применив RoPE Scaling на одной из наших последних наработок – Cotype Pro 8k – нам удалось достичь уровня gpt-4, и даже обогнать ее.

Мы пока не можем выложить приведенные в статье данные и бенчмарки в открытый доступ, однако наши коллеги делятся новым качественным датасетом MTS-AI-SearchSkill/MTSBerquad для задачи генеративного QA!

[1] Ainslie J. et al. Colt5: Faster long-range transformers with conditional computation //arXiv preprint arXiv:2303.09752. – 2023.

[2] Ding J. et al. Longnet: Scaling transformers to 1,000,000,000 tokens //arXiv preprint arXiv:2307.02486. – 2023.

[3] Guo M. et al. LongT5: Efficient text-to-text transformer for long sequences //arXiv preprint arXiv:2112.07916. – 2021.

[4] Chen Y. et al. Longlora: Efficient fine-tuning of long-context large language models //arXiv preprint arXiv:2309.12307. – 2023.

[5] Su J. et al. Roformer: Enhanced transformer with rotary position embedding //Neurocomputing. – 2024. – Т. 568. – С. 127063.

[6] Chen S. et al. Extending context window of large language models via positional interpolation //arXiv preprint arXiv:2306.15595. – 2023.

[7] Peng B. et al. Yarn: Efficient context window extension of large language models //arXiv preprint arXiv:2309.00071. – 2023.