Привет, Хабр!

В прошлом году одним из направлений работы R&D команды в AI VK были энкодеры текстов: модели, которые преобразуют любой текст, от анекдота до официального запроса в техподдержку, в векторное представление — эмбеддинг. Эмбеддинги отражают важные свойства текста, его семантику.

Все энкодеры в NLP можно условно разделить на две группы:

1. Pre‑train‑модели (BERT, RoBERTa, DeBERTa).
   Учатся основным языковым закономерностям, но не умеют явно создавать единый эмбеддинг для всего текста и требуют дообучения под конкретную задачу.

2. Энкодеры текстов (SBERT).
   Сразу выдают готовые семантические эмбеддинги — используются в FAISS, Milvus, других векторных БД. Поверх векторного представления можно применять классические алгоритмы ML. Для оценки схожести текстов просто считаем косинусную близость между эмбеддингами.

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

Для чего нужны энкодеры в NLP?

Наверное, чтобы кодировать. Но как именно что‑то закодировать с помощью классического BERT — вопрос дискуссионный. Формально можно взять эмбеддинг CLS‑токена или использовать любой другой пулинг, но это работает плохо. Поэтому чаще всего BERT используют иначе: поверх него навешивают линейный слой, и модель дообучается под конкретную задачу — например, выдавать один из классов.

Также мы можем оценивать близость с помощью BERT: передаём на вход два текста через сепаратор, учим модель возвращать одно число — оценку близости между текстами. Такие модели есть: их называют cross‑encoder и зачастую они показывают лучшее качество. Проблемы возникают, когда мы хотим масштабировать вычисления. Для ранжирования N текстов придётся выполнить (N^2)/2 сравнений. Кроме того, мы не можем явно хранить векторные представления, а значит, при появлении нового текста нам придётся считать значения близостей между ним и всеми текстами в БД.

Эту проблему решают текстовые энкодеры: они создают один эмбеддинг для всего текста. Полученные векторы можно индексировать, кластеризовать, использовать для поиска ближайших соседей или в качестве признаков. Часто модель можно применять «из коробки», без дообучения — это экономит ресурсы.

ModernBERT

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

Однако в декабре 2024 года была представлена модель ModernBERT — попытка переосмыслить классический BERT с учётом всех современных достижений. Среди ключевых нововведений:

• позиционные эмбеддинги RoPE и поддержка контекста до 8 192 токенов;

• локальное внимание и flash attention;

• оптимизация обучения.

Про эту модель можно почитать в оригинальной статье или этом посте на Habr.

Все отличия от BERT:

ModernBERT прекрасен, но работает только с английским языком. Мы решили это исправить — повторили пайплайн обучения, адаптировав его под русский язык. Так появился RuModernBERT.

Как мы обучали RuModernBERT

Мы полностью воспроизвели трёхэтапную стратегию обучения:

• постепенно увеличивается длина контекста — вплоть до 8 192 токенов;

• уменьшается learning rate;

• данные становятся всё чище.

На первом этапе основную часть обучающей выборки составляет FineWeb (en + ru). Была попытка учить только на нём, но практика показала, что FineWeb не панацея и лучше докидывать других датасетов. Поэтому разбавляем его классическими Wiki, arXiv и т. д. На втором этапе меняем FineWeb на очищенную CulturaX. А на третьем убираем из обучающей выборки и FineWeb, и CulturaX. Параметры learning rate, шедулеров и остального соответствуют обучению оригинальной модели.

Посмотрим, чем RuModernBERT отличается от ModernBERT. Токенизатор учили на русском и английском, основную часть данных составляют русские тексты, размеры моделей смещены в меньшую сторону.

Результаты

Оцениваем модель мы на Russian SuperGLUE — лучше него всё ещё ничего нет. Для каждой задачи мы подбирали гиперпараметры грид‑сёрчем: обучали модель на трейн‑сплите и выбирали лучший результат по валидации.

Сравниваем наши модели с ai‑forever/ruBERT‑base и deepvk/deberta‑v1.

Два последних столбца в таблице — это RuModernBERT-small и RuModernBERT-base. Как видно, модель base стабильно показывает лучшие результаты почти во всех задачах бенчмарка.

Encodechka

Мы не обучали RuModernBERT работать с семантикой, но стало интересно, как он будет взаимодействовать с ней. Для этого мы дополнительно протестировали обе модели на Encodechka — открытом бенчмарке, ориентированном на энкодеры текстов.

Оценка в Encodechka устроена довольно гибко: применяются разные пулинги (mean, CLS), сверху навешиваются логистическая регрессия и KNN, выбирается лучший результат. Это позволяет хотя бы приблизительно оценить потенциал модели в работе с семантикой.

Сравнение проводили с другими претрейнами, результаты которых есть в лидерборде. Итог: RuModernBERT‑base показывает качество, сопоставимое с более крупной ai‑forever/ruRoberta‑large. А версия small немного уступает моделям base.

USER2

USER (Universal Sentence Encoder for Russian) — наше семейство энкодеров текстов. Год назад мы представляли модели USER‑base и USER‑bge‑m3. Когда у нас появился сильный претрейн в лице RuModernBERT, следующий шаг был очевиден — дообучить его для работы с семантикой.

Работ по обучению энкодеров текстов довольно много, и углубляться в особенности каждой из них мы здесь не будем. При обучении USER2 мы решили опираться на пайплайн, предложенный в модели BAAI/bge‑m3, который включает три этапа: RetroMAE, WSFT и SFT.

RetroMAE

Одна из самых актуальных областей применения текстовых энкодеров — задачи retrieval, особенно в контексте RAG (Retrieval-Augmented Generation). Именно поэтому мы включили в пайплайн этап обучения с задачей RetroMAE (Retrieval-oriented Masked Autoencoder), специально разработанной для повышения качества в retrieval-сценариях.

Идея RetroMAE в том, что модель получает на вход текст с высокой долей маскированных токенов, пропускает его через энкодер и через пулинг извлекает эмбеддинг. Затем этот эмбеддинг вместе с ещё более агрессивно маскированной версией того же текста подаётся в небольшой декодер, который должен восстановить исходный текст.

Смысл в том, что декодер восстанавливает текст, опираясь в основном на сжатое векторное представление. Для этого энкодер должен уметь упаковывать всю семантику текста в единый эмбеддинг.

WSFT: много плохих данных

RetroMAE — это отлично, но даже после него эмбеддинги не всегда хорошо сопоставимы по близости. Чтобы исправить это, мы переходим к следующему этапу — WSFT (Weakly Supervised Fine-Tuning).

На этом этапе модель обучается на парах «anchor — positive», где anchor — исходный текст, а positive — текст, близкий к нему по смыслу. В качестве функции потерь используется InfoNCE: для каждой i-й пары положительный пример из пары считается релевантным, а все positive из других пар в батче — нерелевантными (negative). Так модель учится сближать anchor и его positive и одновременно отдалять anchor от всех negatives.

Из-за этого критически важным параметром становится размер батча — чем он больше, тем сложнее задача и тем качественнее обучение.

На WSFT мы больше всего заинтересованы в объёме обучающей выборки. Но для русского языка не так много готовых датасетов с семантически связанными парами, как для других языков.

Поэтому мы применяем две техники:

1. Crosslingual transfer — обучаемся сразу на русском и английском, а также на параллельном корпусе. В этом случае получается значительно нарастить объём — на Hugging Face много таких датасетов.

2. Берём корпус deepvk/cultura_ru_edu и составляем пары на основе простейшей эвристики: из текста выделяем два непересекающихся блока. Суть простая — полученная пара, вероятнее всего, будет связана семантикой. Так мы семплируем 50 млн пар, что уже больше любых отдельных русскоязычных датасетов такого формата.

Как мы уже сказали, для InfoNCE крайне важен размер батча — чем он больше, тем лучше. Благодаря недавней имплементации cached losses в sentence transformers, можно расширять батч практически до бесконечности. Мы остановились на 16 384, как в этой статье, где авторы показывают минимальные улучшения при дальнейшем увеличении батча.

SFT: мало хороших данных

Финальный этап — это SFT (Supervised Fine Tuning). Здесь мы стараемся собрать как можно больше различных качественных датасетов из узких доменов или конкретных задач. Все датасеты мы дополнительно фильтровали с BAAI/bge‑reranker‑v2-m3, отбрасывая 20% пар с наименьшей оценкой близости. Для retrieval‑задач мы также добавляем hard negatives, полученные с помощью USER‑bge‑m3. Всего датасетов 17 штук, с ними можно подробнее ознакомиться в карточках моделей USER2.

В качестве функции потерь взяли GistEmbed. Это улучшенная версия InfoNCE, в которой используется модель‑учитель — у нас multilingual‑e5-small — для динамической фильтрации ложноотрицательных примеров в батче.

Префиксы

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

• “classification: “ — префикс по умолчанию и самый универсальный;

• “clustering: “ — для задач кластеризации: группировка текстов, выявление общих тем, устранение дубликатов;

• “search_query: “ и “search_document: “ — парные префиксы для задач поиска и ранжирования. 

Матрёшка

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

Чтобы на инференсе гибко регулировать необходимую размерность, мы используем метод Matryoshka Representation Learning (MRL). Это подход, позволяющий обучать эмбеддинги так, чтобы их можно было эффективно обрезать по длине. Достигается это пропусканием лосса через несколько усечённых версий эмбеддинга — сначала считаем loss на первых d₁ признаках, затем на d₂, и так далее. В итоге модель обучается формировать представления, где более важная информация сосредоточена в начале вектора.

Мы учили base с размерностями [32, 64, 128, 256, 384, 512 и 768], а small — с [32, 64, 128, 256 и 384].

TLDR этапов обучения USER2

Общие результаты

RuMTEB — это крупнейший русскоязычный бенчмарк для оценки качества энкодеров текстов. Он включает 23 задачи, сгруппированные в 7 категорий. Подробнее про него можно узнать в посте на Хабре. 

Но RuMTEB плохо оценивает работу моделей с длинным контекстом, большая часть задач вмещается в 512 токенов. Поэтому мы отдельно прогнали модель на ru сабсете задачи MLDR (Multilingual Long‑Document Retrieval). MLDR содержит вопросы в паре с очень длинными текстами. Зачастую ответ на вопрос находится далеко за пределами классических 512 токенов, так что можно реально оценить работу длинного контекста.

Для сравнения мы выбрали модели с контекстом 8 192 токена и выше. USER2-base уверенно чувствует себя на уровне более тяжёлых моделей, включая KaLM‑v1.5, Jina‑embeddings и даже E5-mistral-7b.

Мы протестировали, как влияет обрезка выходного эмбеддинга на итоговое качество. Оказалось, что при уменьшении размерности в 3 раза качество представлений практически не меняется!

Итоги

• Обучили RuModernBERT — pretrain‑модель для русского языка на новой архитектуре.

• Построили поверх неё USER2 — sentence encoder в двух версиях: small (34 M) и base (149 M).

• Добавили поддержку MRL, чтобы масштабировать размер эмбеддингов без потерь качества.

• Оценили результат на RSG, Encodechka, RuMTEB и MLDR.

Все модели выложены в open-source, технические детали и примеры использования есть в карточках на Hugging Face:

• USER2-base

• USER2-small

• RuModernBERT-base

• RuModernBERT-small

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