Оптимизация многовекторной эмбеддинговой модели: как пуллинг улучшает качество поиска и экономит место

Многовекторные эмбеддинговые модели предоставляют каждому токену запроса возможность находить соответствия в документе, что делает их крайне выразительными для поиска. Однако количество векторов документов может быстро увеличиваться: корпус с миллионами документов может содержать сотни миллионов векторов! Пуллинг токенов уменьшает эти представления, объединяя похожие векторы токенов. Тем не менее, чем агрессивнее мы сжимаем их, тем легче потерять что-то полезное. Обучение с учетом пуллинга учит кодировщик сохранять важную информацию при объединении векторов, что помогает получать более качественные результаты при том же объеме хранилища. В этой статье мы рассмотрим, как работает пуллинг, обучим модель для сжатого извлечения с помощью Sentence Transformers и исследуем полученные результаты в контексте качества и экономии места.

Почему многовекторное извлечение требует так много векторов

Модель с одним вектором эмбеддинга сжимает весь документ в один вектор. Многовекторные модели, такие как ColBERT, сохраняют коллекцию контекстуальных векторов токенов, позволяя различным частям запроса соответствовать различным частям документа. Например, запрос о ноутбуке с длительным временем работы от аккумулятора может найти информацию о самом ноутбуке в одной части документа, а о времени работы в другой. С помощью MaxSim-оценки каждый вектор запроса выбирает свой лучший соответствующий вектор документа, и эти оценки суммируются: MaxSim(Q,D)=∑q∈Qmaxd∈Dq⊤d Даже с 2-битной или 4-битной остаточной квантизацией, поддерживаемой Fast-PLAID, хранение сотен векторов на документ быстро накапливается. Например, 200 векторов размером 128 измерений при 4 битах на остаточную координату требуют 12.8 КБ квантизированных остатков на документ, или 12.8 ГБ на миллион документов. Пуллинг каждого документа до 50 векторов уменьшает это значение до 3.2 ГБ при той же размерности и точности.

Пуллинг токенов

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

Выбор бюджета пуллинга

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

max_tokens=16 — максимальное количество векторов на документ, включая защищенные токены.

pool_factor=4 — примерно одна четверть оригинальных векторов каждого документа.

distance_threshold=... — общий порог расстояния для объединения, позволяющий количеству векторов варьироваться в зависимости от геометрии эмбеддинга.

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

Калибровка адаптивного пуллинга

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

Обучение с учетом пуллинга

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

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

Я обучил две модели на основе answerdotai/ModernBERT-base, каждая с новым 128-мерным токеном проекции. Одна использовала только кэшированный контрастный убыток, а другая — кэшированный контрастный убыток, обернутый в цели без пуллинга и с фактором 4. Оценка обеих моделей на NanoBEIR@13 показала, что пуллинг-осознанное обучение улучшает качество извлечения при том же объеме векторов.

Заключение

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