Как дообучать мультивекторные модели для специализированного поиска

В статье представлен практический подход к обучению и дообучению моделей многовекторных эмбеддингов с помощью Sentence Transformers. Рабочий процесс разбит на модель, наборы данных, функции потерь, параметры обучения, средства оценки, обратные вызовы и тренер, который связывает все эти компоненты. Автор уделяет особое внимание многовекторному поиску, а не плотным и разреженным эмбеддингам или реранкерам, которым посвящены отдельные руководства. В сопутствующей статье рассматривается эксплуатационная сторона этих моделей, включая их загрузку, кодирование, оценку релевантности и индексацию в векторных базах данных, тогда как эта статья сосредоточена на их адаптации посредством обучения.

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

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

Длина документа представлена как ещё одна критически важная проблема, которая может иметь большее значение, чем выбор архитектуры поиска. Многие общедоступные поисковые модели обучались на фрагментах в стиле MS MARCO и были настроены на относительно короткие входные данные. Классические контрольные точки ColBERT могут обрезать документы после 180 или 300 токенов, а популярные плотные модели обычно ограничиваются 256 или 512 токенами. При применении таких ограничений к более длинным материалам модель незаметно отбрасывает оставшийся текст до вычисления релевантности. Из-за этого она может не увидеть фрагмент, который фактически отвечает на запрос.

В авторской оценке поиска по медицинским материалам используются фрагменты средней длиной 941 токен, а некоторые примеры достигают примерно 1 400 токенов. В таких условиях обрезка снижала NDCG@10 — метрику ранжирования, которая поощряет размещение релевантных результатов ближе к началу выдачи, — на величину до 0,24. Согласно представленному эксперименту, ущерб от неподходящего ограничения длины был значительно больше, чем различия между протестированными архитектурами моделей. Обучение модели для конкретной предметной области позволяет согласовать максимальную длину документа с реальными данными, а не наследовать ограничение, рассчитанное на короткие веб-фрагменты.

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

Таким образом, выбор отправной точки не сводится к поиску модели с лучшим результатом в общедоступных бенчмарках. Специалистам необходимо учитывать, насколько её токенизатор, выученный словарь, соглашения о маркерных токенах, конфигурация проекции, ограничения входных данных и исходная предметная область обучения соответствуют предполагаемой рабочей нагрузке. В качестве примера в статье приводится опыт LightOn в поиске по коду: универсальной модели LateOn оказалось недостаточно для этой специализированной области, поэтому компания обучила LateOn-Code. Та же логика применима к медицинским, юридическим, финансовым, научным и закрытым корпоративным коллекциям, для которых официальная специализированная модель может так никогда и не появиться.

Для полноценной системы обучения также необходимы подходящие данные, функция потерь, преобразующая примеры релевантности в сигнал оптимизации, необязательные параметры обучения для повышения производительности и отслеживания экспериментов, а также средство оценки качества до, во время или после обучения. Класс тренера координирует работу этих компонентов. В статье также рассматриваются использование наборов данных, размещённых на Hugging Face Hub или хранящихся локально, их приведение к требуемому формату, добавление обратных вызовов и обучение на нескольких наборах данных. Эти возможности делают процесс применимым как к простым экспериментам в одной предметной области, так и к более широким смесям обучающих данных.

В качестве конкретной демонстрации при подготовке статьи автор обучил модель multi-vector-encoder/mLateOn-medical. Сообщается, что обучение заняло 14,5 часа на одной Nvidia RTX 3090, а не на крупном кластере ускорителей. В авторской оценке поиска по медицинским материалам полученная модель превзошла все универсальные поисковые системы, которые удалось найти, включая плотные, разреженные, лексические и многовекторные альтернативы. Практический вывод состоит в том, что тщательно настроенная и адаптированная к предметной области многовекторная модель способна за несколько часов превзойти универсальные поисковые системы на потребительском оборудовании, особенно если она сохраняет полную длину и специализированную лексику целевых документов.