Привет, Хабр! Меня зовут Саприн Семён. Я занимаюсь анализом данных и машинным обучением в компании ПГК Диджитал. Сегодня мы продолжаем серию статей, в которой я рассказываю о том, как мы с командой разрабатывали ИИ-помощника. В прошлой статье мы обсудили, почему стандартные подходы к работе с документами не всегда работают, и какие шаги помогли нам повысить качество поиска без существенных затрат памяти на GPU.

Сегодня речь пойдёт о следующем этапе: дообучении (fine-tuning) модели эмбеддингов для улучшения качества поиска в RAG-системе. Это позволило нам получить более точные представления документов и пользовательских запросов, что напрямую сказалось на релевантности финальных ответов. Давайте перейдём к деталям.

Описание проблемы

Как вы уже знаете из первой части, RAG-система состоит из двух ключевых компонентов: ретривера и генератора.

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

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

Сегодня существует множество мощных предобученных моделей, таких как deepvk/USER-bge-m3, которые показывают отличные результаты на общих задачах и доменах. Однако при работе с узкоспециализированными данными — в нашем случае это юридические документы — такие модели не всегда способны адекватно оценить схожесть между запросом и чанком.

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

Идея обучения

Мы решили проводить обучение подобно алгоритму, предложенному NVIDIA: NV-Retriever: Improving text embedding models with effective hard-negative mining.

Основной проблемой создания эмбеддингов на таких узких доменах является тот факт, что векторные представления релевантных и нерелевантных чанков по косинусной мере будут близки друг к другу, и, как следствие, к вопросу пользователя (если он в принципе релевантен к данным). Поэтому в качестве функции потерь (loss function) наш выбор пал на triplet-margin-loss, чтобы помочь модели научиться более строго различать релевантные и нерелевантные чанки.

Для этого нам необходимо сформировать обучающую выборку следующего вида:

1.    Anchor – запрос пользователя
   Например: «Какие существуют способы согласования договорной цены?»

2. Positive context – релевантный контекст (чанк, в котором действительно приведён ответ на вопрос пользователя)
   Например: «1.2 Способы согласования договорной цены ...»

3. Negative context  - нерелевантный контекст (любой другой чанк, в котором нет ответа на вопрос)
   Например: «1.4 Способы закупки…»

Чтобы улучшить точность обученной модели, мы использовали hard-negative примеры вместо обычных negative.

Hard-negative примеры – это такие чанки, которые базовая модель может ошибочно посчитать за положительные. Другими словами, это сложные случаи, в которых модель с трудом может определить разницу между релевантным и нерелевантным контекстом. В рассмотренном выше примере в качестве Hard-negative context мы хотим получить что-то вроде «1.3 Регламент изменения договорной цены …».

Суть в следующем: мы даем модели триплет (вопрос, положительный контекст, отрицательный контекст) и говорим:

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

Просмотрев множество таких примеров, модель подстроится под домен и в будущем станет более точно отличать релевантный чанк от псевдо-релевантного, тем самым полученный ретривером контекст будет более точный и полный, что увеличит качество ответов генератора. 

Перейдём к первому шагу – формирование выборки.

Формирование обучающей выборки

Как рассказывал в прошлой статье, у нас есть сформированный экспертами небольшой датасет ~100 экземпляров вида «вопрос + чанк». В нашем случае это «anchor + positive».

Также коллеги предоставили хороший корпус юридических документов. Выбираем из него несколько общих документов и получаем чанки ~1000 штук (чанки берём согласно методу из прошлой части статьи). Далее необходимо сформировать триплеты. Первым делом перейдём к созданию вопросов с помощью LLM (в нашем случае взяли 8-ми битный квант Qwen2.5-14B).

С использованием готового функционала (generate_qa_embedding_pairs от llama-index), проходим по полученным чанкам и создаём 10 вопросов для каждого из них. Тем самым мы получим ~10.000 пар anchor-positive, где anchor – сгенерированный вопрос, а positive – релевантный чанк.

Пример промпта для создания вопроса:

PROMPT = """
«Вы - ИИ помощник, которому поручено сгенерировать реалистичные пары вопрос-ответ на основе заданного документа». Вопрос должен быть таким, который пользователь мог бы задать при поиске информации, содержащейся в документе.
Дано: {context_str}
Инструкции:
1. Проанализируй ключевые темы, факты и концепции в данном документе, выберите ту, на которой стоит сосредоточиться.
2. Сформулируй {num_questions_per_chunk} похожих вопросов, которые пользователь мог бы задать, чтобы найти в этом документе информацию, не содержащую названия компании.
3. Используй естественный язык и иногда добавляй в вопрос опечатки или разговорные выражения, чтобы имитировать реальное поведение пользователя, но не указывай на их наличие.
4. Убедитесь, что вопрос семантически связан с содержимым документа, без прямого копирования фраз.
5. Убедитесь, что все вопросы похожи друг на друга. т.е. все они касаются схожей темы/запрашивают одну и ту же информацию, но разными способами.
6. Сгенерируй вопросы, которые максимально похожи на запросы обычного человека. Используй неформальные фразы и выражения.
Прояви творческий подход, подумай, как любознательный пользователь и сгенерируй {num_questions_per_chunk} похожих вопросов, которые приведут к данному документу при семантическом поиске.
Ниже приведи только {num_questions_per_chunk} вопросов, каждый в новой строке.
"""

Может показаться, что создание аж 10 вопросов по каждому чанку – избыточно, но в нашем случае это происходит с учетом того, что модель будет ошибаться и не всегда отрабатывать идеально, тем самым необходимо фильтровать полученные результаты, ведь сгенерированные вопросы могут быть нерелевантные к чанку / слишком общими и не содержать конкретики / быть не на русском языке и т.д.

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

Пример промпта для фильтрации пар anchor-positive:

PROMPT = """
Твоя задача - распознать неудачные пары вопрос / контекст.
Неудачными считаются следующие пары:
1) Вопрос написан не на русском языке
2) На вопрос невозможно ответить, основываясь только на информации из контекста
3) Контекст содержит неинформативные данные
4) Вопрос плохо сформулирован или является неполным
Ниже приведены пары в формате:
JSON: [("query_ids": id пары, "anchor": вопрос, "positive": контекст)]
Дано:
{inputs}
В качестве ответа верни только id тех пар, которые являются неудачными. Не нужно приводить обоснование, или свои комментарии.
"""

Таким образом у нас получилось оставить ~3000 пар (из 10.000 сформированных на первом шаге). Сохраняем отфильтрованные пары и идём дальше.

Hard-negative mining

Здесь мы как раз и формируем итоговые триплеты. Как я и писал раньше, в качестве негативов можно взять случайные чанки из полного набора, за исключением тех, что являются позитивными для конкретного примера, но мы идем дальше и хотим сформировать hard-negative. Для этого нам потребуются, так называемые, модели – учителя, а также базовая модель (deepvk/USER-bge-m3), которую мы в дальнейшем будем обучать.

Алгоритм создания hard-negative примеров:

1. Предобработка данных:
   1)     Для каждого элемента обучающего датасета (anchor-positive) мы сохраняем уникальный идентификатор для удобного поиска в дальнейшем.

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

2. Генерация hard-negative примеров для одной пары anchor-positive:
   1)     Векторизуем anchor (текущий вопрос) и соответствующий ему positive (положительный чанк).

   2)     Вычисляем косинусное сходство между anchor и positive— это становится "базовым" порогом (base similarity).

   3)     Задаем целевой порог отбора: threshold * base_similarity, где threshold — гиперпараметр.

   4) Считаем косинусную меру между каждым positive других пар и текущим anchor и ранжируем их по уменьшению косинусной меры. Далее выбираем те, у которых сходство с anchor меньше заданного на предыдущем шаге порога, и берём n самых близких к нему (но всё ещё нерелевантных) чанков.

   Мы провели исследование, сформировав несколько датасетов с разными параметрами threshold и n, и пришли к выводу, что лучшие результаты итоговой модели получаются с параметрами: threshold=0.97 n=3

3. Использование нескольких моделей-учителей
   1)     Чтобы повысить разнообразие и качество hard-negative примеров, мы используем несколько моделей-учителей (3 в нашем случае):
   ·       intfloat/multilingual-e5-large-instruct
   ·       sergeyzh/rubert-tiny-turbo
   ·       cointegrated/LaBSE-en-ru

   2)     Для каждой модели повторяем описанный выше процесс, собирая hard-negative кандидатов для каждой пары anchor-positive

   3) Результирующий hard-negative пример формируется как случайный выбор из объединённого множества кандидатов, собранных всеми моделями-учителями.

4. Финальная сборка датасета
   После того как для каждого запроса собраны hard-negative примеры, мы формируем финальный триплет вида (anchor, positive, negative) и добавляем его в обучающий датасет.

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

Обучение модели

1. Модель и её настройка
   
   Мы используем предобученную модель deepvk/USER-bge-m3 как базовую архитектуру. Чтобы минимизировать использование памяти GPU, мы применяем LoRA-адаптеры, которые позволяют обучать только небольшое количество дополнительных параметров поверх уже обученной модели.

   Наша конфигурация LoRA:
   ·       r=16 — ранг матрицы адаптеров,
   ·       lora_alpha=32 — коэффициент масштабирования,
   ·       target_modules=['query', 'key', 'value', 'dense'] — блоки добавления адаптеров,
   ·       lora_dropout=0.1 —dropout для регуляризации.

   trainable params: 7,110,656 | all params: 366,137,344 |trainable%: 1.9421

   Такая настройка позволила нам сохранить около 98% параметров модели необучаемыми, что существенно снизило требования к памяти и ускорило обучение.

   Рекомендации для настройки гиперпараметров LoRA адаптера:

   1)     Задайте базовые значения для коэффициента масштабирования адаптера lora-alpha и ранга матриц r
   2)     Настройте скорость обучения learning rate
   3)     После нахождения learning rate настройте ранг матриц r
   4)     Нет необходимости изменять значение lora_alpha (вы можете продолжать использовать значение с первого шага)

2. Датасеты
   
   Трейн и валидационный датасеты получены путём перемешивания итогового датасета, сформированного в пункте 4 предыдущего параграфа. Соотношение размеров трейн к валидации – 4:1. Трейн датасет содержит около 2400 триплетов.

   Тестовый датасет – те самые ~100 пар anchor-positive, полученные при помощи экспертов, с прикрученными к ним hard-negative по указанному выше алгоритму.

3. Лосс
   Стоит упомянуть о выбранном лоссе – TripleMarginLoss. Ниже приведена его математика:
   
   — эмбеддинги anchor, positive и negative соответственно,
   margin — гиперпараметр, задающий "зазор" между классами.

   Всё, как и говорил ранее. Цель: сделать так, чтобы расстояние между anchor и positive было меньше чем anchor и negative хотя бы на margin.

4. Метрики качества
   Для оценки качества использовались две основные метрики:
   1)     NDCG@K — учитывает порядок релевантности документов, применяя экспоненциальное взвешивание к первым позициям.
   2)     Recall@K — доля релевантных документов среди первых K результатов.

   Финальная модель сохраняется, если значение NDCG@5 на валидации выше предыдущего лучшего. В наших экспериментах взяли параметр К=5, чтобы в дальнейшем (при инференсе RAG системы) не раздувать контекст до 10 чанков, а стремиться к тому, чтобы релевантные чанки приходили в контекст как можно выше.

   Если у читателя возник вопрос: «Как вы считаете метрики ранжирования на валидации, если в каждом вашем примере только один релевантный и один нерелевантный документ?» – вы совершенно правы. Мы поступаем следующим образом:
   1)     Для каждого запроса (anchor) мы:
   ·       Вычисляем эмбеддинг;
   ·       Считаем эмбеддинги всех чанков из датасета (а не только с одним positive и одним negative);
   ·       Получаем полный список косинусных сходств.
   2)     Далее мы:
   ·       Сортируем чанки по убыванию сходства с запросом;
   ·       Смотрим, на какой позиции находится нужный (релевантный) чанк;
   3) На основании этого рассчитываем NDCG@K и Recall@K.
   

5. Обучение
   Далее всё как обычно, всем известный цикл обучения с AdamW, Cosine Scheduler with WarmUp и отображением метрик после каждой эпохи. 40 эпох. 2 часа на A100 (40 gb) и смотрим к чему мы пришли.

Результаты

После завершения обучения мы провели оценку качества нашей дообученной модели эмбеддингов на тестовом датасете. Мы сравнили показатели базовой модели (до обучения) и нашей версии с LoRA-адаптерами после fine-tuning. Вот как выглядят результаты:

Все метрики показали значительный прирост после дообучения. Например:

Recall@5 вырос более чем на 10%, то есть на пятом месте поисковой выдачи наша модель находит релевантный документ в 79% случаев, против 67% у исходной версии.

Заключение

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

• Дообучению на доменных данных, специфичных для нашей предметной области;

• Использованию hard-negative примеров, которые помогли модели лучше различать похожие, но нерелевантные чанки;

• Применению LoRA, которое позволило эффективно обучать модель без значительного расхода GPU-ресурсов.

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

P.S. Если остались вопросы по алгоритму или коду (который я здесь не стал приводить), можете прийти ко мне в личку (tg: @huraligne).