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

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

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

• MassiveIntentClassification

• MassiveScenarioClassification

• MIRACL Reranking

• MIRACL Retrieval

• STS22

• RUParaphraserSTS

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Данные

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

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

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

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

Обучение

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

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

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

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

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

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

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

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

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

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

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

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

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

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

За основу мы взяли опубликованные датасеты из работ bge-m3 (данные на Hugging Face)  и Nomic Embed (конфиг данных), которые предназначены как раз для этой стадии обучения. Однако мы внесли ряд изменений: провели фильтрацию негативных примеров, отобрали лишь часть датасетов. Например, убрали 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)

Анна Максимова (@anpalmak)

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

Ссылки

FRIDA

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

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