Обучение с учетом объединения токенов для многовекторных моделей эмбеддингов

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

Объединение токенов (Token pooling) уменьшает размер этих представлений путем слияния похожих векторов токенов. Однако чем агрессивнее я их сжимаю, тем легче потерять что-то полезное. Обучение с учетом объединения (Pooling-aware training) учит энкодер сохранять важную информацию при слиянии векторов, помогая улучшить качество поиска при том же объеме хранилища.

В этом посте я объясню, как работает объединение (pooling), обучу модель для поиска в сжатом виде с помощью Sentence Transformers и исследую компромисс между качеством и объемом хранилища. При pool_factor=4 моя обычно обученная модель сохраняет 94,6% своего исходного показателя NanoBEIR nDCG@10 при 4-кратном сокращении количества векторов документов. Обучение с учетом объединения токенов повышает этот показатель до 97,8% при том же объеме хранилища! Давайте посмотрим, как это работает.

Этот пост основан на обучении с учетом объединения, изученном в работе Стефана Йозефа Learn to Pool и исследовании регуляризации объединения от LightOn, в котором также рассматриваются целевое обучение и адаптивные бюджеты. Здесь я сосредоточусь на понимании метода и его применении на коротком воспроизводимом примере с использованием Sentence Transformers.

Этот пост основан на обучении с учетом объединения, изученном в работе Стефана Йозефа Learn to Pool и исследовании регуляризации объединения от LightOn, в котором также рассматриваются целевое обучение и адаптивные бюджеты. Здесь я сосредоточусь на понимании метода и его применении на коротком воспроизводимом примере с использованием Sentence Transformers.

Обучение с учетом объединения улучшает качество поиска при том же бюджете векторов. Подробности ниже.

Оглавление

• Почему многовекторному поиску нужно так много векторов
• Объединение токенов Выбор бюджета объединения Калибровка адаптивного объединения
• Выбор бюджета объединения
• Калибровка адаптивного объединения
• Обучение с учетом объединения токенов Обертка функции потерь в Sentence Transformers Опциональное ускорение ядра
• Обертка функции потерь в Sentence Transformers
• Опциональное ускорение ядра
• Использование обученной модели
• Результаты оценки
• Комбинирование методов сжатия
• Ссылки
• Код и модели

Почему многовекторному поиску нужно так много векторов

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

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

MaxSim⁡(Q,D)=∑q∈Qmax⁡d∈Dq⊤d \operatorname{MaxSim}(Q, D) = \sum_{q \in Q} \max_{d \in D} q^\top d

Даже с учетом 2-битного или 4-битного остаточного квантования, поддерживаемого Fast-PLAID, хранение сотен векторов на документ обходится дорого. При 4 битах на остаточную координату 200 векторов размерностью 128 требуют 12,8 КБ квантованных остатков на документ, или 12,8 ГБ на миллион документов. Объединение (pooling) каждого документа до 50 векторов сокращает этот объем до 3,2 ГБ при той же размерности и точности.

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

Объединение токенов

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

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

Вы можете сделать сжатие пропорциональным длине документа:

При pool_factor=4 документы сохраняют примерно четверть своих исходных векторов, поэтому более длинные документы могут сохранять пропорционально большую емкость. По умолчанию первый токен защищен от слияния, а объединение применяется только к эмбеддингам документов, оставляя эмбеддинги запросов без изменений.

Выбор бюджета объединения

HierarchicalTokenPooling предлагает три альтернативных способа управления:

Это различие важно при разной длине документов. С pool_factor=4 более длинный документ может сохранить примерно четверть своих векторов, что позволяет его представлению расти вместе с содержанием. С distance_threshold количество векторов может увеличиваться, когда необходимо сохранить больше различных эмбеддингов.

Калибровка адаптивного объединения

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

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

Для калибровки требуются только эмбеддинги документов; она возвращает порог, который можно повторно использовать при индексации. Цель — среднее значение, поэтому отдельные документы могут получить больше или меньше 16 векторов. Достигнутое среднее значение также может измениться на другом корпусе.

Обучение с учетом объединения токенов

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

Если вы использовали эмбеддинги Matryoshka, принцип обучения покажется вам знакомым. Matryoshka применяет цель поиска к нескольким вариантам ширины эмбеддинга, обучая меньшие префиксы оставаться полезными. Обучение с учетом объединения применяет цель к представлениям с меньшим количеством векторов токенов.

В моем примере каждый тренировочный батч вносит вклад в две функции потерь:

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

Я усредняю их:

L=12Lunpooled+12Lpooled \mathcal{L} = \tfrac{1}{2}\mathcal{L}_{\mathrm{unpooled}} + \tfrac{1}{2}\mathcal{L}_{\mathrm{pooled}}

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

Кластеризация Уорда принимает дискретные решения о том, какие токены должны быть объединены. Во время обучения HierarchicalTokenPooling использует прямой оценочный метод (straight-through estimator) для передачи градиентов от объединенных векторов обратно к их исходным токенам без дифференцирования назначений кластеров. Таким образом, энкодер может обучаться на ошибках поиска, допущенных при использовании его объединенных представлений.

Обертка функции потерь в Sentence Transformers

Обертка охватывает существующую функцию потерь для многовекторного поиска и применяет ее с использованием выбранных вами пулеров:

CachedMultiVectorMultipleNegativesRankingLoss обеспечивает целевую функцию поиска. Она обучает каждый запрос оценивать релевантный документ выше, чем негативные документы в пакете. По умолчанию используется model.similarity, применяющая настройки сходства модели (в данном примере MaxSim) и настройки ядра. Модель, сконфигурированная с similarity_fn_name="meanmaxsim", также будет использовать MeanMaxSim во время обучения.

Кэширование градиентов позволяет энкодеру обрабатывать небольшие мини-пакеты, сохраняя при этом больший набор негативных примеров внутри пакета. Мой скрипт обучения использует мини-пакеты по 64 при общем размере обучающего пакета 256.

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

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

Замените пулер с коэффициентом 4 на выбранную вами конфигурацию развертывания или включите несколько пулеров и выберите веса для их потерь.

Полный скрипт обучения включает в себя набор данных, трейнер, функцию потерь при валидации и небольшие оценки поиска во время обучения. После клонирования репозитория запустите оба варианта локально с помощью uv:

Скрипт объявляет свои зависимости для обучения и ядер, поэтому uv устанавливает их за вас. Совместимые ядра включаются автоматически.

Вы также можете проводить обучение на облачном GPU с помощью Hugging Face Jobs. Установите CLI и войдите в систему с токеном, который позволяет запускать задания и загружать модели:

Это запускает вариант с учетом пулинга на GPU A10G с ограничением по времени в четыре часа. Передача HF_TOKEN в качестве секрета позволяет скрипту загрузить обученную модель в вашу учетную запись Hub после завершения обучения. Чтобы также запустить базовый вариант с обычной функцией потерь, повторите команду с флагом --baseline после URL-адреса скрипта.

Дополнительное ускорение с помощью ядер

Вы можете ускорить оценку MaxSim, кластеризацию Уорда и поддерживаемые операции бэкенда с помощью kernels. Установите зависимости для обучения и ядер:

В приведенном выше примере обучения замените инициализацию модели и пулера перед созданием функции потерь:

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

Sentence Transformers выбирает совместимые ядра для вашей настройки и использует существующую реализацию везде, где замена недоступна. Например, ядро MaxSim по умолчанию оценивает запросы относительно документов, не создавая полную матрицу сходства токенов в памяти, но для этого требуются CUDA и Triton.

Ядро кластеризации Уорда по умолчанию значительно ускоряет иерархическую кластеризацию за счет работы на CUDA вместо CPU. Речь идет об ускорении в 30-150 раз по сравнению с SciPy, в зависимости от настроек.

Использование обученной модели

Давайте загрузим модель с поддержкой пулинга и выполним поиск со сжатыми эмбеддингами документов. Настройки ядер не сохраняются вместе с моделью, поэтому добавьте use_kernels=True при загрузке, если хотите использовать их и для инференса:

Пулер сжимает эмбеддинги документов после конвейера прямого прохода энкодера, сокращая общее количество с 58 векторов до 16, оставляя эмбеддинги запросов без изменений. Mars по-прежнему получает наивысшую оценку сходства, за ним следуют Jupiter и Venus, поэтому ранжирование сохраняется, несмотря на хранение гораздо меньшего количества векторов документов.

Вы можете изменить pool_factor при инференсе, чтобы изучить другой бюджет хранения с той же обученной моделью. Для повторного использования вы также можете передать token_pooling=pooling при инициализации MultiVectorEncoder, чтобы сделать это значением по умолчанию во время выполнения. Задачи пулера по-прежнему контролируют, какие входные данные сжимаются (например, только документы).

Результаты оценки

Для исследования я обучил две модели на базе answerdotai/ModernBERT-base, каждая с новой 128-мерной проекцией токенов:

• tomaarsen/mve-ModernBERT-base-token-pooling-baseline, использующая только кэшированную контрастивную функцию потерь.
• tomaarsen/mve-ModernBERT-base-token-pooling, использующая кэшированную контрастивную функцию потерь, обернутую в указанные выше целевые функции для непулированных данных и данных с коэффициентом 4.

Оба запуска используют одну эпоху на одних и тех же 120 000 триплетах MS MARCO, seed 42, а также идентичные настройки инициализации и оптимизации. Ограничение на количество токенов при обучении составляет 300, а короткие запросы расширяются как минимум до 32 позиций токенов для предоставления дополнительных контекстных векторов запроса. Разница между двумя экспериментами заключается только в функции потерь, и я не проводил поиск гиперпараметров.

Я оценил обе модели на NanoBEIR@13, сообщив среднее значение nDCG@10 (%). На рисунке сравнивается обычное обучение и обучение с учетом пулинга с коэффициентами пулинга 2, 3, 4 и 6, наряду с непулированным поиском. В таблице также представлены фиксированные ограничения и адаптивные пороги, показывающие, как преимущества обучения переносятся на другие настройки пулинга.

Для адаптивного пулинга я откалибровал отдельный порог для каждого энкодера и корпуса NanoBEIR, используя его непулированные эмбеддинги документов, чтобы соответствовать среднему количеству векторов при ограничениях 64, 32 и 16. Короткие документы снижают эти средние значения ниже соответствующих ограничений.

При обученном pool_factor=4 показатель nDCG@10 улучшается с 49.03 до 51.23, при этом хранится на 74.8% меньше векторов документов, чем при непулированном поиске. Это увеличивает сохранение оценки с 94.6% до 97.8%, где сохранение — это среднее значение nDCG@10 после пулинга, деленное на собственное среднее значение непулированной модели. Непулированный поиск также улучшается с 51.84 до 52.39.

Прирост распространяется на фиксированные ограничения и адаптивный пулинг, даже несмотря на то, что при обучении использовались только полное представление и пропорциональный пулинг. При среднем значении 29.58 векторов на документ адаптивный пулинг сохраняет 97.3% от непулированной оценки модели с учетом пулинга, храня при этом на 84.7% меньше векторов документов. Обычная обученная модель сохраняет 93.3% при том же бюджете.

Комбинирование методов сжатия

Если размер индекса является вашей главной заботой, я бы обратился к этим вариантам в следующем порядке:

• Квантование эмбеддингов до 4, 2 или 1 бита, например, через , чтобы уменьшить объем памяти на координату.
• Пулинг токенов для уменьшения количества векторов документов, опционально с дообучением с учетом пулинга, как показано в этом блоге.
• Эмбеддинги токенов Matryoshka для уменьшения размерности каждого вектора токена, опционально с дообучением для сохранения качества поиска при меньшей ширине.

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

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

Попробуйте полный пример обучения, выберите бюджет, который важен для вашего приложения, и сравните его с соответствующим запуском с обычной функцией потерь. Сопутствующий скрипт оценки измеряет извлечение без пулинга, ограничения в 64, 32, 16, 8 и 4 вектора, а также коэффициенты пулинга от 2 до 6, чтобы вы могли построить график компромисса между качеством и объемом хранения для вашей собственной модели.

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

Ссылки

• Штефан Йозеф. Learn to Pool: Lightweight Fine-Tuning for Flexible Multi-Vector Compression.
• Бенджамин Клави, Антуан Шаффен и Гриффин Адамс. Сокращение объема многовекторного извлечения с минимальным влиянием на производительность с помощью пулинга токенов.
• Антуан Шаффен. После вечеринки наступает бесплатный обед: регуляризация моделей ColBERT для улучшения возможностей пулинга и уменьшения размера индекса.
• Луюй Гао, Юньи Чжан, Цзявэй Хань и Джейми Каллан. Масштабирование размера пакета глубокого контрастивного обучения в условиях ограниченной памяти.
• Введение в модели встраивания Matryoshka.

Код и модели

• Пример обучения с учетом пулинга токенов и README
• Полный скрипт обучения
• Скрипт оценки NanoBEIR
• Примеры контрольных точек: tomaarsen/mve-ModernBERT-base-token-pooling tomaarsen/mve-ModernBERT-base-token-pooling-baseline
• tomaarsen/mve-ModernBERT-base-token-pooling
• tomaarsen/mve-ModernBERT-base-token-pooling-baseline
• Предварительно обученные многовекторные модели, совместимые с Sentence Transformers
• Обучение многовекторных моделей с учетом Matryoshka