Библиотека Sentence Transformers выпустила руководство по обучению ColBERT-моделей эмбеддингов

Том Аарсен показывает в блоге Hugging Face, как обучать и дообучать мультивекторные модели эмбеддингов, так называемые ColBERT-style модели с «поздним взаимодействием» (late interaction), с помощью класса MultiVectorEncoder библиотеки Sentence Transformers. В отличие от обычной, «плотной» (dense) модели эмбеддингов, которая сжимает весь текст в один вектор, мультивекторная модель хранит по отдельному вектору на каждый токен и сравнивает запрос с документом через операцию MaxSim: для каждого токена запроса ищется его наиболее похожий токен в документе, и такие максимумы суммируются. Сопоставление на уровне токенов сохраняет тонкие детали, которые единый вектор вынужден усреднять, за счёт более крупного индекса модель точнее находит релевантные документы.

Дообучение под конкретный домен даёт заметный прирост, потому что словарь, стиль запросов и само понятие релевантности различаются между веб-поиском, юридическим поиском, поиском по коду и научной литературой: сопоставление токен к токену позволяет мультивекторным моделям улавливать именно эти доменные тонкости, и даже небольшое количество доменных примеров заметно улучшает результат. Отдельная проблема готовых моделей, длина документа: классические чекпоинты ColBERT обрезают документы на 180 или 300 токенах, многие популярные плотные модели, на 256 или 512, а среди опубликованных мультивекторных чекпоинтов типичный потолок, от 180 до 512 токенов. На медицинском тестовом наборе автора, где пассажи занимают в среднем 941 токен, такое усечение стоит до 0,24 по метрике точности поиска NDCG@10, это значительно больше, чем разница между самими архитектурами моделей. Похожая история случилась у компании LightOn: её универсальная модель LateOn не справлялась с поиском по коду, и компании пришлось обучить отдельную модель LateOn-Code.

Центральный эксперимент поста, сравнение шести стартовых точек для дообучения. Автор взял шесть разных чекпоинтов, дообучил каждый по идентичному рецепту на 25 тысячах пар «вопрос, пассаж» из медицинского датасета MIRIAD и оценил результат на 1000 отложенных вопросах против корпуса из 50 тысяч пассажей. Итог оказался неожиданным и повторился в двух семействах моделей: чекпоинты уровня «unsupervised», прошедшие крупномасштабное контрастное предобучение, но ещё не общую supervised-настройку под универсальный поиск, дообучаются под новый домен значительно лучше своих уже финализированных версий и даже обгоняют их, хотя стартуют с более низкого результата; финализированные же чекпоинты почти не сдвигались с места или проседали при любой опробованной скорости обучения. Как более слабую, но рабочую альтернативу автор проверил и «холодный старт»: случайно инициализированная проекционная голова поверх сильного претренированного бэкбона Alibaba-NLP/gte-modernbert-base на тех же 25 тысячах пар почти сравнялась с готовыми чекпоинтами, разрыв составил всего 0,03.

В качестве рабочего примера всей методики автор дообучил модель multi-vector-encoder/mLateOn-medical на полном датасете MIRIAD (4,4 миллиона пар «медицинский вопрос, пассаж», точное число строк, 4 467 542) за 14,5 часа на одной потребительской видеокарте RTX 3090. По его данным, эта модель легко обошла все проверенные им универсальные модели поиска, плотные, разреженные (sparse), лексические и мультивекторные.

Из практических приёмов пост фиксирует: у части чекпоинтов (пример, GTE-ModernColBERT-v1 с лимитами query_length=48 и document_length=300) ограничения длины запроса и документа можно снять вручную под свои более длинные тексты, у автора медицинские пассажи доходят до 1400 токенов, а семейство mLateOn, например, изначально работает на полном контексте бэкбона в 8192 токена; при сборке модели с нуля на базовом трансформере (в примере, answerdotai/ModernBERT-base) библиотека сама добавляет случайно инициализированную токенную проекцию (128 измерений в примере), которую затем обучает класс MultiVectorEncoderTrainer; включение скиплиста пунктуации, исключающего знаки препинания из индексации документа, в 4-вариантной проверке (без скиплиста / пунктуация / стоп-слова / оба) немного повысило качество и заодно бесплатно сократило размер индекса документов на 9,6% на данных автора; а тестирование классического приёма ColBERT, расширения запроса токенами [MASK], в четырёх конфигурациях не дало измеримой разницы, так что автор не считает его обязательным. Материал продолжает серию постов автора о дообучении эмбеддинг-моделей (плотных, разреженных, реранкеров) и связан с отдельным постом о том, как уже готовые мультивекторные модели загружать, кодировать и индексировать в векторных базах.

Ключевые факты

Sentence Transformers получила класс MultiVectorEncoder для обучения и дообучения мультивекторных (ColBERT-style) моделей эмбеддингов, «с нуля» на базовом трансформере или поверх готового чекпоинта.
На примере медицинского поиска: модель multi-vector-encoder/mLateOn-medical, дообученная за 14,5 часа на одной RTX 3090 на 4,4 миллиона пар «вопрос, пассаж» из MIRIAD, по данным автора легко обошла все проверенные им универсальные модели поиска, плотные, разреженные, лексические и мультивекторные.
Сравнение шести стартовых чекпоинтов на 25 тысячах пар из MIRIAD показало: «unsupervised»-версии моделей дообучаются под домен заметно лучше уже финализированных версий и даже обгоняют их, хотя стартуют ниже; финализированные чекпоинты почти не сдвигались или проседали.
Многие готовые чекпоинты обрезают документы в диапазоне 180-512 токенов; на медицинских пассажах (в среднем 941 токен в датасете MIRIAD) такое усечение стоит до 0,24 NDCG@10, значительно больше, чем разница между архитектурами.
Практические приёмы: скиплист пунктуации в оценке документа немного повышает качество и бесплатно сокращает индекс на 9,6%; классическое [MASK]-расширение запроса в четырёх протестированных конфигурациях не дало измеримой разницы.

Почему это важно

Мультивекторные (ColBERT-style) модели эмбеддингов сопоставляют текст токен к токену через операцию MaxSim, а не сжимают его в один вектор, это сохраняет тонкие сигналы релевантности, которые «плотная» (dense) модель вынуждена усреднять, и обычно даёт более точный поиск ценой более крупного индекса. Пост показывает, что готовые модели такого типа настроены на короткие фрагменты (многие чекпоинты обрезают документ в диапазоне 180-512 токенов) и на общий, а не доменный поиск, а дообучение под свой домен на модели, прошедшей контрастное предобучение, но ещё не финальную supervised-настройку, даёт максимальный прирост. Это подробный, воспроизводимый рецепт такого дообучения внутри библиотеки Sentence Transformers, доведённый до рабочего примера на одной потребительской видеокарте за 14,5 часа.

Кому это важно

ML/NLP-инженерам и командам, которые строят поиск или retrieval-часть RAG-системы под конкретный домен, медицину, юриспруденцию, финансы, код или внутренние документы компании, и которым не хватает точности готовых универсальных моделей эмбеддингов. Полезно и тем, кто выбирает между dense-, sparse- и мультивекторным поиском для длинных документов: пост прямо показывает цену обрезки длинных текстов готовыми моделями. Отдельно материал важен командам без крупной GPU-инфраструктуры: весь пример воспроизведён на одной RTX 3090.

Как это применить

В библиотеке Sentence Transformers для этого есть класс MultiVectorEncoder. Первый путь, дообучить готовый мультивекторный чекпоинт (в примере, lightonai/mLateOn-unsupervised): токены-маркеры запроса и документа, проекционная голова и скиплист остаются от чекпоинта, менять стоит прежде всего лимиты длины запроса и документа под свои тексты (query_length и document_length, вплоть до None, если чекпоинт этого не требует). Второй путь, взять произвольный базовый трансформер (в примере, answerdotai/ModernBERT-base): библиотека сама добавит случайно инициализированную токенную проекцию (128 измерений в примере), которую дообучает класс MultiVectorEncoderTrainer практически с нуля. Данные, простые пары «запрос, релевантный пассаж»: их можно взять готовыми с Hugging Face Datasets Hub (в посте, 4,4-миллионный медицинский датасет MIRIAD) или собрать локально из CSV, JSON, Parquet, Arrow или SQL. Выбор стартовой точки, по опыту автора: сначала искать «unsupervised»-версию модели своего семейства; нет такой, строить свежую проекцию на сильном ретривал-бэкбоне (разрыв с готовыми чекпоинтами в тесте автора, всего 0,03); дообучать уже финализированный универсальный чекпоинт, по его данным, самый слабый из трёх вариантов. Из необязательных, но проверенных настроек, скиплист пунктуации в оценке документа, который немного повышает качество и сокращает индекс, и отказ от классического [MASK]-расширения запроса, если оно не даёт прироста на своих данных.

Можно ли доверять

Это разбор от первого лица автора самого инструмента, а не независимый бенчмарк, но методика описана подробно и проверяемо: контролируемое сравнение шести стартовых чекпоинтов по одинаковому рецепту на одинаковых данных (25 тысяч пар из MIRIAD, проверка на 1000 отложенных вопросах против корпуса 50 тысяч пассажей), это весомее голословного заявления. При этом главный результат флагманского примера, что дообученная модель mLateOn-medical «легко обошла» все протестированные автором универсальные модели поиска, подкреплён цифрами: точный показатель NDCG@10 у mLateOn-medical, 0,9139 против 0,8520 у лучшего универсального конкурента lightonai/mLateOn (отрыв +0,062), а всё сравнение проведено и оценено самим автором на собственном тестовом наборе, без независимой проверки третьей стороной.

Риски и подводные камни

Все цифры примера, усечение стоит до 0,24 NDCG@10, скиплист экономит 9,6% индекса, разрыв «холодного старта» в 0,03, получены на одном домене (медицинские тексты MIRIAD) и одной связке модель+данные конкретного автора; переносить их как гарантию на другой домен, язык или объём данных нельзя без повторного замера. Мультивекторный подход по определению хранит вектор на каждый токен, а не один вектор на документ, индекс получается заметно крупнее, чем у dense-модели, а во сколько именно раз он вырастет на реальных объёмах, в материале не оценивается. Модель mLateOn-medical и её веса выложены в открытый доступ на Hugging Face Hub, готовый чекпоинт для своего медицинского поиска читатель получает, дообучать самому не придётся. Наконец, время дообучения (14,5 часа на одной RTX 3090) не переведено в стоимость, во сколько обойдётся аналогичный прогон в облаке, из поста не следует.

««Unsupervised»-чекпоинты дообучаются под новый домен куда лучше своих финализированных версий, и даже обгоняют их, хотя стартуют с более низкого результата.»

— Том Аарсен