Знакомьтесь, FRIDA. Открытая эмбеддинг-модель для русского языка

В этой статье мы расскажем о нашей новой модели FRIDA, которая сейчас (20.05.2025) занимает первое место в русскоязычном бенчмарке MTEB (ссылка на таблицу лидеров ^[1]).

Ранее мы уже рассказывали ^[2] на Хабре о создании русскоязычных задач для MTEB. Напомним, что этот бенчмарк предназначен для оценки моделей, способных создавать эмбеддинги текста — векторные представления, применяемые в различных задачах NLP.

ruMTEB, русскоязычная часть MTEB, включает в себя 23 задания. Это 17 уникальных русскоязычных датасетов, собранных нашими коллегами, и шесть мультиязычных, взятых из оригинального MTEB:

• MassiveIntentClassification

• MassiveScenarioClassification

• MIRACL Reranking

• MIRACL Retrieval

• STS22

• RUParaphraserSTS

Создавая модель FRIDA, мы применили подходы, которые ранее успешно использовали при разработке модели ru-en-RoSBERTa ^[3], а также добавили несколько новых приёмов, про часть из которых мы расскажем в этой статье.

Подготовка базовой модели

Архитектура нашей модели основывается на энкодерной части FRED-T5-1.7B ^[4] (Full-scale Russian Enhanced Denoisers based on T5 architecture). В библиотеке HuggingFace эта архитектура реализована в виде T5EncoderModel. FRED-T5 был изначально обучен на данных на русском языке. По версии бенчмарка Russian SuperGLUE, оценивающего общую способность к пониманию естественного языка, производительность модели наравне ^[5] с Mistral 7B.

Сейчас открытых данных для обучения ^[6] эмбеддинг-модели на русском языке немного, поэтому мы используем преимущественно данные на английском языке, полагаясь на трансфер знаний между языками. Чтобы усилить способности модели в понимании и обработке английского языка, мы обновили токенизатор. Он был создан путём добавления английских токенов из оригинальной RoBERTa в токенизатор от FRED-T5.

После расширения изначального словаря возникла необходимость проинициализировать матрицу эмбеддингов. Оригинальные эмбеддинги из FRED-T5 мы оставили без изменений. Для новых токенов из RoBERTa применили следующий подход: токены проходили токенизацию через исходный словарь FRED-T5, затем мы получали эмбеддинги этих токенов и усредняли их. Так формировался эмбеддинг для нового токена. Этот метод показал лучшие результаты по сравнению с рандомной инициализацией и частичной инициализацией на основе эмбеддингов RoBERTa.

Так как у FRIDA появились новые токены в словаре для английского языка, а FRED-T5 во время предобучения видел данные преимущественно на русском языке, мы сначала немного адаптируем модель. Для этого мы её обучаем с помощью стандартной задачи для энкодер-моделей — Masked Language Modeling (MLM). 

В качестве данных мы использовали уникальные тексты из того же русско-английского датасета, что и при разработке ru-en-RoSBERTa (таблица 6 в статье ^[7]). Датасет покрывает большинство интересных для нас доменов, видов задач и единиц текста — это предложения, параграфы и текст целиком. Всего получается около 11 млн текстов.

Превращение в эмбеддинг-модель 

Для создания эмбеддинг-модели обычно применяют подход к обучению в два этапа. Сначала предобучают (contrastive pre-training) базовую модель на большом объёме неразмеченных примеров — например, 200 млн. Затем дообучают (contrastive fine-tuning) на сравнительно небольшом объёме качественных, размеченных примеров — около 1 млн. 

Для небольших моделей уровня BERT-large этап контрастивного предобучения даёт существенный прирост на задачах, связанных с поиском. Однако эксперименты показывают ^[8], что моделям покрупнее, например, Mistral 7B, такой этап не требуется. Это связывают с масштабным предобучением (casual language modeling) подобных моделей на триллионах токенов. Тем не менее, мы решили проверить подход с двумя этапами обучения в действии, потому что энкодер FRED-T5 всего в два раза больше BERT-large. 

Инструкции и префиксы

Использование инструкций в обучающих примерах пришло в эмбеддинги из обучения LLM. Инструкции помогают модели лучше понимать, какие пары текстов считать релевантными (далее – пример из этой ^[9] статьи). Действительно, для запроса пользователя «Реализация батч-нормализации в Python» релевантными могут оказаться самые разные документы:

• «Я хочу разработать батч-нормализацию с нуля», то есть дубликат вопроса;

• «Вы можете просто импортировать torch.nn.BatchNorm2d» — непосредственно ответ на вопрос;

• Просто фрагмент кода.

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

Этот подход должен помогать модели обобщаться на новые задачи, а также учитывать специфику задачи, которая закладывается в инструкцию. На деле ряд работ (раз ^[10], два ^[11]) показывает, что модели скорее используют инструкции как набор ключевых слов, чем обращают внимание ^[12] на детали описания. Поэтому мы используем вместо инструкций префиксы, которые сообщают, какой класс задачи предстоит решать.

Мы взяли за основу схему префиксов, предложенную командой Cohere ^[13], которая использовалась и в ru-en-RoSBERTa, но немного изменили её. У нас получилось семь префиксов для конкретных задач:

1. «search_query:» и «search_document:» предназначены для поиска ответа или документа, который может содержать ответ;

2. «paraphrase:» — для задач, связанных с перефразированием, в основном на уровне предложений (STS, дедупликация);

3. «categorize:» – для асимметричного сопоставления заголовка и тела документа (например, новостей, научных статей, постов в социальных сетях); 

4. «categorize_sentiment:» предназначен для любых задач, которые полагаются на фичи, связанные с эмоциями ^[14] и прочим;

5. «categorize_topic:» — для задач, где нужно группировать тексты по темам;

6. «categorize_entailment:» – для логических задач по типу, в основном NLI.

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

Контрастивное предобучение

Задача этого этапа — научить модель находить релевантный запросу документ среди множества других. Для построения обучающих пар «запрос—документ» обычно не используют ручную разметку — достаточно естественной структуры текстовых данных. Например, можно взять заголовок статьи как запрос, а её содержание — как документ. Или вопрос с форума как запрос, а соответствующий ответ — как документ. Также возможны пары из автоматически переведённых текстов. Они позволяют обучать модель выделять смысловое сходство между запросом и документом даже без явной разметки людьми.

Данные

Для русского языка мы собрали коллекцию ^[15] датасетов из открытых источников, которую назвали solyanka. В совокупности получилось 10 миллионов пар. Каждый датасет мы дедуплицировали и отобрали тексты по длине и качеству. Фильтрация была выполнена с использованием метода consistency filtering из статьи E5 ^[16]. Его суть в следующем:

1. Для каждого запроса определяются ближайшие документы из всего корпуса и запоминается, на каком месте (обозначим его R) в списке ближайших соседей (k-nearest neighbors) находится документ относительно своего запроса.  Мы использовали FAISS для построения индекса поиска.

2. В выборку включались только те примеры, где значение R<N. В оригинальной статье для каждого датасета N=2. Мы обнаружили, что значение N лучше подбирать индивидуально для каждого набора данных.

Напомним, что для этапа предобучения обычно используют на порядок больше пар. Например, 200 млн при работе над упомянутой E5 и около 3 млрд в mGTE ^[17]. Поэтому мы добавили часть данных, которые опубликовали авторы Nomic Embed (конфиг данных ^[18]). Эти данные представляют собой примеры на английском языке, которые также были очищены с помощью метода consistency filtering. Для эксперимента мы взяли 10% от всего датасета, оставив только те, что содержат не более 512 токенов. В итоге мы получили русско-английский корпус на 30 млн примеров. 

Обучение

Во время контрастивного предобучения батч формируется из нескольких пар запрос-документ. Мы штрафуем модель за неверное определение релевантного документа исходя из близости между запросами и документами. 

Так как мы используем одновременно несколько датасетов, то может случиться так, что в батче окажется более одного релевантного к запросу документа. Поэтому мы собираем в батч данные только из одного датасета. Это исключает появление ложных негативных примеров (false negatives) из других датасетов и обеспечивает единообразие распределения данных внутри батча.

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

Обмен негативными примерами между GPU

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

Использование Grad-Cache

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

1. Сначала считываются и сохраняются эмбеддинги запросов и документов без построения графа вычислений.

2. Далее мы получаем так называемый Representation Gradient Cache, когда по микробатчам считаем лосс и градиент только для полученных в пункте 1 эмбеддингов.

3. Повторяем ^[19] вычисления для каждого микробатча, но уже строим для эмбеддинга граф вычислений и считаем градиенты. Таким образом, мы итеративно накапливаем полный градиент.

Этот подход помог эффективно увеличить размер батча, не выходя за рамки доступной памяти ^[20] GPU. Обучение проводилось на 8 GPU H100.

Дообучение на качественных данных

Для этого этапа (contrastive fine-tune) мы использовали специализированные датасеты с относительно высоким качеством, однако общее число примеров невелико и насчитывает всего 1 млн. 

Мы выбирали датасет в основном по двум критериям: полезен ли датасет для решения практических задач? И можем ли мы оценить, как эти задачи решаются? Например, мы использовали датасет MS MARCO для задачи поиска по веб-документам. Практически задача нам интересна, но метрики мы можем посчитать только для английского языка. Однако этот датасет полезен для поиска в целом и, в частности, для задачи поиска по Википедии, которую мы уже можем оценить благодаря MIRACL и RuBQ.

За основу мы взяли опубликованные датасеты из работ bge-m3 (данные ^[21] на Hugging Face)  и Nomic Embed (конфиг данных ^[22]), которые предназначены как раз для этой стадии обучения. Однако мы внесли ряд изменений: провели фильтрацию негативных примеров, отобрали лишь часть датасетов. Например, убрали QA-датасеты на тему права и медицины, потому что эти задачи оценить не можем. Дополнительно включили примеры из разных доменов на классификацию и кластеризацию. Пары сформировали следующим образом: в качестве запросов используются случайные N текстов. Позитивами будут считаться тексты того же класса, а негативами — тексты из других классов. Мы используем примеры только из трейн сплитов и всегда проверяем пересечения с тестом.

Отдельно хочется упомянуть фильтрацию негативных примеров и в целом их майнинг. Мы воспользовались следующей идеей: для каждого запроса считается близость для всех его документов (позитивных и негативных). Затем считаем порог отбора как среднее значение близости по позитивным документам и умножаем на коэффициент, например, равный 0,95. Таким образом, для фильтрации или майнинга отбирается top-k документов, удовлетворяющих заданному порогу, который специфичен для каждого запроса.

Финальная полировка модели проходит в две стадии. Сначала мы обучаемся только на задачах поиска и используем как hard-negative документы, так и in-batch negative. Затем добавляем все оставшиеся данные для categorize и paraphrase префиксов и обучаемся дальше, но с использованием только hard-negative документов. Такой подход позволяет добавить  в обучение данные для категоризации, для которых in-batch стратегия приведёт к тому, что негативными примерами окажутся документы из того же класса.

В итоге нам удалось существенно улучшить результаты предыдущей модели ru-en-RoSBERTa (+9.2), а также обойти «инструктивные» модели E5. Например, при обучении только на одном contrastive fine-tune датасете в две стадии FRIDA получила выше 70 пунктов на всём бенчмарке — во многом благодаря сильному энкодеру от FRED-T5 и этапу подготовки данных. Также важным для задач поиска оказался этап contrastive pre-training (+1.8 в среднем). Возможно, на большем масштабе (100+ млн пар) мы увидели бы и больший прирост, однако для модели 800M эффект от предобучения уже не ярко себя проявляет.

Заключение

В этой статье мы представили новую модель FRIDA, а также набор данных для разработки и исследования текстовых эмбеддинг-моделей. Мы также рассмотрели основные подходы и результаты, которые они показали.

Далее мы планируем развивать способы оценки эмбеддингов и пробовать новые подходы к их разработке. И обязательно делиться своими наработками с сообществом!  

Авторы

Команда RnD AI, ML для B2C: 

Артем Снегирев (@artemsnegirev ^[23])

Анна Максимова (@anpalmak ^[24])

Александр Абрамов (@Andriljo ^[25])

Ссылки

FRIDA ^[26]

Лидерборд ^[1] (выбрать MTEB(rus, v1))

Датасеты для предобучения ^[15]