Обучение и дообучение многовекторных моделей эмбеддингов с помощью Sentence Transformers

Обновить на GitHub

Sentence Transformers — это библиотека Python для использования и обучения моделей эмбеддингов и переранжирования в широком спектре приложений, включая генерацию с дополнением поиска, семантический поиск, семантическое текстовое сходство и многое другое. Обновление v6.0 добавляет четвёртый тип модели: MultiVectorEncoder для поиска с поздним взаимодействием в стиле ColBERT, а также полный подход к её обучению. В этой статье я покажу, как дообучить многовекторную модель, которая превзойдёт универсальные модели поиска на ваших данных. Этот метод также позволяет обучать с нуля новые эффективные многовекторные модели. Всё ниже работает после выполнения pip install -U "sentence-transformers[train]".

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

Наконец, в разделе «Оценка» я покажу, что моя дообученная модель multi-vector-encoder/mLateOn-medical, обученная за 14,5 часа на одной RTX 3090 одновременно с подготовкой этой статьи, легко превосходит все найденные мной универсальные модели поиска в оценке поиска по медицинским данным: плотные, разреженные, лексические и многовекторные.

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

Эта статья посвящена обучению многовекторных моделей. Если вы хотите узнать, как их использовать — от загрузки и кодирования до индексации в векторных базах данных, — ознакомьтесь с сопутствующей статьёй «Многовекторные модели эмбеддингов (позднее взаимодействие) с Sentence Transformers».

Содержание

• Что такое многовекторные модели?
• Зачем выполнять дообучение?
• Компоненты обучения
• Модель
  • Дообучение существующей многовекторной модели
  • Создание модели на основе базового трансформера
  • Какую исходную точку выбрать?
• Набор данных
  • Данные на Hugging Face Hub
  • Локальные данные
  • Формат набора данных
• Функция потерь
• Параметры обучения
• Оценщик
• Тренер
  • Обратные вызовы
  • Обучение на нескольких наборах данных
• Оценка
  • Оптимизация индекса
• Благодарности
• Дополнительные ресурсы
  • Примеры обучения
  • Документация

Что такое многовекторные модели?

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

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

Зачем выполнять дообучение?

Дообучение многовекторных моделей значительно повышает качество поиска в конкретной предметной области: словарь, стиль запросов и понятие релевантности различаются для веб-поиска, юридического анализа, поиска по коду и обзора научной литературы. Поскольку запросы и документы сопоставляются токен за токеном, многовекторные модели улавливают тонкие доменные сигналы, которые модели с одним вектором обычно усредняют, и очень хорошо реагируют даже на умеренный объём внутридомен­ных данных для дообучения.

Кроме того, большинство выпущенных моделей поиска рассчитано на короткие фрагменты. Классические контрольные точки ColBERT обрезают документы до 180 или 300 токенов, а многие популярные плотные модели — до 256 или 512, поскольку их обучающие данные в стиле MS MARCO редко выходят за эти пределы. Если ваши документы длинные, такие модели незаметно отбрасывают большую часть каждого документа до его оценки. В моей медицинской оценке, где средняя длина фрагмента составляла 941 токен, я измерил, что это усечение снижает NDCG@10 максимум на 0,24 — значительно сильнее, чем различия между архитектурами моделей. Обучая собственную модель, вы можете задать длину документа, необходимую именно вашим данным.

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

Компоненты обучения

Обучение моделей MultiVectorEncoder включает следующие компоненты:

1. Модель: модель для дообучения или архитектура для создания с нуля.
2. Набор данных: данные, используемые для обучения и оценки.
3. Функция потерь: функция, измеряющая качество модели и направляющая процесс оптимизации.
4. Параметры обучения (необязательно): параметры, влияющие на производительность обучения, отслеживание и отладку.
5. Оценщик (необязательно): класс для оценки модели до, во время или после обучения.
6. Тренер: объединяет все компоненты обучения.

Рассмотрим каждый компонент подробнее.

Модель

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

Дообучение существующей многовекторной модели

Если вы хотите дополнительно дообучить существующую многовекторную модель, вам вообще не нужно беспокоиться об архитектуре:

from sentence_transformers import MultiVectorEncoder


model = MultiVectorEncoder(
    "lightonai/mLateOn-unsupervised",
    model_kwargs={"torch_dtype": "float32"},
    processor_kwargs={"model_max_length": 8192},  
)

Контрольная точка содержит собственную конфигурацию: маркеры запросов и документов, проекционную голову и список исключаемых токенов. При дообучении обычно следует сохранить всё это и менять только то, что требует специфика ваших данных. Сначала проверьте конфигурацию длины, поскольку многие опубликованные контрольные точки ограничивают документы 180–512 токенами (см. раздел «Зачем выполнять дообучение?»), тогда как мои медицинские фрагменты достигают 1400 токенов. Семейство mLateOn уже использует полный контекст базовой модели длиной 8192 токена, но если исходная контрольная точка имеет ограничения, снимите их:

model[0].query_length = None
model[0].document_length = None

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

Я внёс ещё одно изменение: добавил список исключаемых знаков препинания, чтобы токены пунктуации не учитывались при оценивании и хранении документов. В абляционном исследовании с четырьмя вариантами (без списка, пунктуация, стоп-слова, оба варианта) это умеренно улучшило качество, а индекс документов бесплатно уменьшился на 9,6% на этих данных:

import string


model[2].skiplist_words = list(string.punctuation)
model[2].resolve_with_tokenizer(model.tokenizer)  

Создание модели на основе базового трансформера

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

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("answerdotai/ModernBERT-base", model_kwargs={"torch_dtype": "float32"})

Это классический конвейер ColBERT: Transformer создаёт контекстуализированные эмбеддинги токенов, Dense на уровне токенов проецирует каждый из них в пространство размерности 128, MultiVectorMask определяет, какие токены учитывать при оценивании, а затем применяется нормализация на уровне токенов. Проекция начинается со случайных значений, поэтому до обучения эта модель непригодна для использования. Интересно, что этот подход работает и с сильными базовыми моделями плотных эмбеддингов. Новая проекция на основе Alibaba-NLP/gte-modernbert-base в моих экспериментах отставала от исходных контрольных точек всего на 0,03, используя только проекцию и 25 тысяч обучающих пар.

Классические приёмы токенизации ColBERT ([MASK]-расширение запроса, префиксные токены [Q] / [D], ограничение длины документа и список исключаемой пунктуации) по умолчанию отключены и настраиваются. Полный набор описан в разделе «Создание пользовательских моделей». Я проверил [MASK]-расширение запроса в четырёх конфигурациях для дообучения в своей области, и ни одна не дала измеримого улучшения, поэтому не обязательно использовать классическую схему.

Какую исходную точку выбрать?

Я измерил это непосредственно при подготовке статьи: взял шесть исходных точек и обучил каждую по одной и той же схеме на 25 тысячах пар «медицинский вопрос — фрагмент» из MIRIAD, а затем оценил на 1000 отложенных вопросах относительно корпуса из 50 тысяч фрагментов:

Результат меня удивил и повторился в двух семействах моделей. Контрольные точки с суффиксом -unsupervised значительно лучше адаптируются к новой области, чем их завершённые аналоги, и, несмотря на более низкий стартовый результат, обходят их. Эти контрольные точки находятся после масштабного контрастивного предварительного обучения, но до контролируемого дообучения для общего поиска: они сохраняют всю структуру позднего взаимодействия, но не имеют универсальной настройки, которую затем приходится отменять при обучении в конкретной области. Завершённые контрольные точки, напротив, почти не изменились или даже ухудшились при всех проверенных мной скоростях обучения.

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

Набор данных

MultiVectorEncoderTrainer использует экземпляры datasets.Dataset или datasets.DatasetDict для обучения и оценки. Данные можно загрузить из Hugging Face Datasets Hub или использовать локальные данные в любом удобном формате (например, CSV, JSON, Parquet, Arrow или SQL).

Примечание: множество общедоступных наборов данных, работающих с Sentence Transformers из коробки, помечены на Hugging Face Hub тегом sentence-transformers, поэтому их легко найти по адресу https://huggingface.co/datasets?other=sentence-transformers. Просмотрите их, чтобы найти готовые наборы данных, полезные для ваших задач, областей или языков.

Данные на Hugging Face Hub

Для загрузки данных из наборов на Hub можно использовать функцию load_dataset:

from datasets import load_dataset

train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train")

print(train_dataset)
"""
Dataset({
    features: ['question', 'passage_text'],
    num_rows: 4467542
})
"""

Именно на этом наборе данных я буду обучать модель в статье: 4,4 миллиона медицинских вопросов из MIRIAD, каждый из которых сопоставлен с исходным фрагментом, содержащим ответ (в среднем 941 токен). Простые пары «запрос — релевантный фрагмент» — самые лёгкие для сбора обучающие данные для собственной области, и, как вы увидите, их достаточно.

Локальные данные

Для загрузки локальных данных в распространённых форматах также можно использовать load_dataset:

from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_file.csv")

dataset = load_dataset("json", data_files="my_file.json")

Если локальные данные требуют предварительной обработки, для инициализации набора данных из словаря списков можно использовать datasets.Dataset.from_dict:

from datasets import Dataset

queries = []
documents = []



dataset = Dataset.from_dict({
    "query": queries,
    "document": documents,
})

Формат набора данных

Важно, чтобы формат набора данных соответствовал функции потерь (или чтобы функция потерь соответствовала формату данных). Проверка совместимости выполняется в два этапа:

1. Если согласно таблице «Обзор функций потерь» функции потерь требуется метка, набор данных должен содержать столбец с именем «label» или «score». Этот столбец автоматически используется как метка.
2. Все столбцы, не названные «label» или «score», согласно таблице «Обзор функций потерь» считаются входными данными. Число оставшихся столбцов должно соответствовать числу допустимых входов выбранной функции потерь. Имена этих столбцов не имеют значения, важен только порядок.

Помимо этого, для многовекторных моделей действуют два специальных соглашения:

• Позиционное назначение запроса и документа: первый столбец кодируется как запрос, а все последующие — как документы, независимо от имён столбцов. Это поведение можно переопределить для каждого столбца стандартным параметром обучения router_mapping.
• Формат дистилляции знаний: по одному столбцу на каждый документ-кандидат, то есть (query, document_1, ..., document_N, scores), где scores — список из N оценок учителя для каждой строки. Для наборов данных KD, где идентификаторы запросов и документов хранятся отдельно от текстовых наборов (например, lightonai/ms-marco-en-bge), можно использовать resolve_ids, чтобы разрешать идентификаторы в тексты на лету.

Функция потерь

Функции потерь количественно оценивают качество модели на заданной пачке данных, позволяя оптимизатору обновлять веса модели и получать более благоприятные (то есть меньшие) значения потерь. Выбор функции потерь зависит от доступных данных и целей. Полный список вариантов приведён в разделе «Обзор функций потерь».

Для распространённых пар «вопрос — ответ» или «вопрос — фрагмент» основным вариантом является обучение с внутрипакетными отрицательными примерами с помощью MultiVectorMultipleNegativesRankingLoss, где каждый другой документ в пакете служит отрицательным примером для каждого запроса. Большие пакеты означают больше отрицательных примеров и более эффективное обучение, поэтому на практике стоит использовать вариант GradCache — CachedMultiVectorMultipleNegativesRankingLoss, который отделяет эффективный размер пакета от объёма, помещающегося в памяти GPU:

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.losses import CachedMultiVectorMultipleNegativesRankingLoss

model = MultiVectorEncoder("lightonai/mLateOn-unsupervised", model_kwargs={"torch_dtype": "float32"})

loss = CachedMultiVectorMultipleNegativesRankingLoss(
    model=model,
    mini_batch_size=16,  
)

Параметр mini_batch_size ограничивает использование памяти, кодируя документы порциями такого размера, тогда как эффективный контрастивный размер пакета (128 в моём запуске; в абляционных экспериментах увеличение пакета не дало дополнительной пользы) остаётся независимым. GradCache гарантирует одинаковые результаты независимо от размера порции, поэтому на меньших GPU его можно уменьшить ценой только увеличения времени выполнения. Если длина документов сильно различается, рассмотрите параметр mini_batch_num_tokens, который формирует каждую порцию по общему бюджету токенов, а не по числу документов. Так необычно длинные документы не смогут резко увеличить потребление памяти (мой mini_batch_size=16 при примерно 940 токенах на документ соответствует mini_batch_num_tokens=15_000).

Одна специфическая для многовекторных моделей ловушка заключается в том, что контрастивные функции потерь по умолчанию используют scale=1.0, в отличие от эквивалентного варианта для плотных эмбеддингов, где по умолчанию задано scale=20.0. Значение 20.0 необходимо потому, что косинусное сходство — это одно число в диапазоне [-1, 1], слишком узком для резкого softmax. Оценка MaxSim вместо этого суммирует одно сходство лучшего совпадения для каждого токена запроса и уже находится примерно в диапазоне [0, query_length]: 32-токенный запрос может получить оценку до 32. Поэтому не переносите scale=20.0 из скрипта обучения плотной модели: это приведёт к насыщению softmax и обнулению градиентов.

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