Многомерные модели эмбеддингов, такие как ColBERT, обеспечивают высокую точность поиска, но количество хранимых векторов на документ быстро увеличивает стоимость индекса. Техника группировки, токен-пулинг, уменьшает этот размер, и, согласно источнику, адаптированное обучение позволяет лучше сохранять качество результатов при равном бюджете хранения. При pool_factor=4 модель, обученная обычным способом, сохраняет 94,6% своего несжатого nDCG@10 на NanoBEIR, против 97,8% для модели, обученной с учётом пулинга.
Почему многомерные модели хранят так много векторов
Классическая модель эмбеддингов сводит весь документ к одному вектору. Многомерные модели, напротив, сохраняют набор контекстных векторов, по одному на токен (единицу текста, часто слово или часть слова), что позволяет разным частям запроса соответствовать разным частям документа. В качестве примера источник приводит запрос о ноутбуке с длительным временем автономной работы, который может найти упоминание об устройстве в одной части документа и упоминание о времени работы батареи в другой.
Такая работа имеет свою цену. Согласно представленному расчёту, при квантовании 4 бита на остаточную координату, 200 векторов размерности 128 представляют 12,8 КБ на документ, то есть 12,8 ГБ для миллиона документов. Сокращение каждого документа до 50 векторов уменьшает этот объём до 3,2 ГБ при той же размерности и точности.
Метод основан на наблюдении: многие векторы одного документа могут быть близки, поскольку несколько токенов часто описывают схожие понятия. HierarchicalTokenPooling, библиотека Sentence Transformers, группирует эти векторы с помощью кластеризации Уорда, которая на каждом шаге объединяет две группы, вызывающие наименьшее увеличение внутреннего разброса. Затем каждая группа заменяется своим средним значением и нормализуется. При pool_factor=4 документ сохраняет примерно четверть своих исходных векторов, так что более длинный документ пропорционально сохраняет больше векторов.
Обучить модель сжатию
Чем сильнее сжатие, тем легче потерять полезную информацию. Обучение с учётом пулинга как раз направлено на сохранение важной информации при объединении векторов. Каждая обучающая партия сочетает две функции потерь: поиск с исходными эмбеддингами документов и поиск после сжатия примерно до четверти векторов. Эти две функции потерь усредняются, и обе цели обновляют один и тот же энкодер.
Поскольку кластеризация принимает дискретные решения, источник указывает, что обучение использует straight-through оценщик, который передаёт градиенты от сгруппированных векторов к составляющим их токенам, без дифференциации назначения кластерам. Класс MultiVectorTokenPoolingLoss инкапсулирует существующую функцию потерь для поиска, в данном случае CachedMultiVectorMultipleNegativesRankingLoss, с весами 0,5 для каждого варианта. Исходная модель — answerdotai/ModernBERT-base, обученная на 120 000 триплетов MS MARCO в течение одной эпохи.
Результаты на NanoBEIR
При pool_factor=4 nDCG@10 увеличивается с 49,03 до 51,23 для модели, обученной с пулингом, одновременно сокращая количество векторов документов на 74,8% по сравнению с несжатым поиском. Несжатая модель также улучшается с 51,84 до 52,39. При среднем значении 29,58 векторов на документ адаптивный пулинг сохраняет 97,3% несжатого результата обученной модели, против 93,3% для модели, обученной обычным способом.
Для уменьшения размера индекса источник предлагает сначала квантование эмбеддингов, затем токен-пулинг, возможно, с тонкой настройкой с учётом пулинга, и, наконец, матрёшечные эмбеддинги (Matryoshka embeddings) для уменьшения размерности каждого вектора. Эти подходы можно комбинировать. Работы, на которых основана статья, описаны в исследовании Learn to Pool Стефана Йозефа.
