Token Pooling-Aware Training: сжатие векторов без потери качества

Проблема масштабируемости мульти-векторных моделей

Мульти-векторные модели (например, ColBERT) обеспечивают высокую точность поиска, сохраняя контекстный вектор для каждого токена документа. Однако это создает проблему хранения: корпус из миллиона документов может требовать сотен миллионов векторов. При использовании 4-битной квантования 200 векторов размерности 128 занимают 12.8 КБ на документ, что в масштабах миллиона документов выливается в 12.8 ГБ только остаточных данных.

Решение: Token Pooling и иерархическая кластеризация

Метод Token Pooling уменьшает размер представлений путем слияния схожих векторов токенов с помощью иерархической кластеризации Уорда (Ward hierarchical clustering). Алгоритм последовательно объединяет группы с минимальным увеличением дисперсии внутри кластеров. Это позволяет гибко контролировать бюджет хранения:

• pool_factor=4: документы сохраняют ~25% исходных векторов.
• max_tokens=16: жесткое ограничение на количество векторов.
• distance_threshold: адаптивное сжатие на основе геометрии эмбеддингов.

Обучение с учетом пулинга (Pooling-Aware Training)

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

L = 0.5 * L_unpooled + 0.5 * L_pooled

Это заставляет энкодер формировать такие токенические эмбеддинги, чтобы полезные для поиска признаки сохранялись даже после агрегации векторов. Градиенты передаются через "straight-through estimator", позволяя модели учиться на ошибках поиска в сжатом пространстве.

Результаты и метрики

Эксперименты на датасете NanoBEIR (метрика nDCG@10) показали значительное преимущество предложенного метода по сравнению с обычным обучением при одинаковом бюджете хранения (pool_factor=4):

Метод интегрирован в библиотеку Sentence Transformers через модуль HierarchicalTokenPooling, что позволяет воспроизвести результаты и внедрить оптимизацию в существующие пайплайны RAG-систем.

Источник: Hugging Face blog ↗