ColBERT

ColBERT (контекстуализированное позднее взаимодействие на основе BERT) — нейросетевая модель поиска, которая представляет запрос и каждый документ одним эмбеддингом на токен и оценивает документ с помощью позднего взаимодействия: для каждого токена запроса находится наиболее похожий токен документа, а максимальные значения сходства, называемые MaxSim, суммируются в оценку релевантности. Omar Khattab и Matei Zaharia из Стэнфордского университета представили эту модель в статье SIGIR 2020 года ColBERT: эффективный и результативный поиск фрагментов с помощью контекстуализированного позднего взаимодействия на основе BERT. Эмбеддинги токенов документов вычисляются заранее, офлайн, поэтому во время запроса выполняется недорогое вычисление сходства, а не полный проход BERT для каждого документа-кандидата.

Роль позднего взаимодействия#

Нейросетевые поисковые модели различаются тем, когда запрос и документ взаимодействуют. Биэнкодеры сжимают каждый текст в один вектор — это стандартная схема для векторного поиска: векторы документов вычисляются заранее, и поиск выполняется быстро, но один вектор сглаживает различия между отдельными терминами. Кросс-энкодеры одновременно обрабатывают запрос и документ с помощью трансформера, поэтому механизм внимания сопоставляет каждую пару токенов. Они обеспечивают наивысшую точность и должны запускаться для каждой пары «запрос — документ», поэтому подходят только для переранжирования коротких списков кандидатов.

ColBERT кодирует запрос и документ отдельно, как биэнкодер, но сохраняет по одному 128-мерному вектору на токен и позволяет двум сторонам взаимодействовать только на заключительном этапе MaxSim. Запросы дополняются маскирующим токеном BERT до фиксированной длины 32, благодаря чему у модели появляются дополнительные позиции для расширения запроса связанными терминами. В тестах поиска фрагментов MS MARCO и TREC CAR ColBERT показал результаты на уровне ранжировщиков на основе BERT, при этом работал в сто раз быстрее и использовал на четыре порядка меньше FLOPs на запрос.

Вычисление оценки MaxSim#

Оператор MaxSim реализуется всего в несколько строк PyTorch. В этом примере он применяется к контекстуальным эмбеддингам токенов, полученным с помощью небольшого энкодера семейства BERT:

import torch.nn.functional as F
from sentence_transformers import SentenceTransformer

encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
query = "who wrote the novel Dune"
docs = [
    "Frank Herbert published the science fiction novel Dune in 1965.",
    "A dune is a hill of sand shaped by the wind.",
]

# One L2-normalized contextual embedding per query token
q = F.normalize(encoder.encode(query, output_value="token_embeddings"), dim=-1)
for doc in docs:
    d = F.normalize(encoder.encode(doc, output_value="token_embeddings"), dim=-1)

    # MaxSim: each query token keeps its best-matching document token, and the maxima are summed
    score = (q @ d.T).max(dim=1).values.sum()
    print(f"{score.item():.2f}  {doc}")

Оценка для фрагмента о романе составляет 5.90, а для фрагмента о песчаных дюнах — 4.33. В обоих есть слово «dune», которое хорошо совпадает в каждом случае, но в первом фрагменте токен запроса «wrote» находит «published» со сходством по косинусу 0.72, а «novel» находит сам себя; во втором фрагменте наилучшее совпадение для «wrote» составляет всего 0.29.

ColBERT и близкие термины#

• Сам BERT: предварительно обученный энкодер, который ColBERT дообучает для поиска, добавляя линейный слой, проецирующий каждый токен в пространство размерности 128.
• Позднее разбиение на фрагменты: Текст тоже кодируется с учетом контекста до разделения, а затем токены усредняются в один вектор на фрагмент; ColBERT же сохраняет вектор каждого токена.
• Семантический поиск: Задача, которую решает ColBERT, ранжируя фрагменты по смыслу, а не по совпадающим ключевым словам.
• Гибридный поиск: Объединяет результаты поиска по ключевым словам и по одному вектору, чтобы совместить точное совпадение терминов со смысловым поиском. Сопоставление токенов в ColBERT обеспечивает часть такой точности в рамках одной нейросетевой модели.

Позднее взаимодействие для изображений#

Позднее взаимодействие применимо и к изображениям. ColPali (2024) кодирует снимок страницы документа как сетку векторов патчей с помощью визуально-языковой модели и оценивает его относительно токенов запроса с помощью MaxSim, перенося позднее взаимодействие в мультимодальный RAG для PDF-файлов. По схожей причине детекция объектов с открытым словарём сохраняет детализацию признаков изображения: Ultralytics YOLOE хранит текстовый эмбеддинг для каждого заданного класса, например «двухэтажный автобус», и сравнивает его с признаками областей во всех точках карт признаков, поэтому короткая подсказка может соответствовать одному небольшому объекту, а не одному вектору, суммирующему всё изображение.

ColBERTv2 и PLAID#

Хранение вектора для каждого токена делает индекс ColBERT примерно на порядок больше индекса с одним вектором. Эти затраты важны, когда ColBERT используется как поисковая модель первого этапа или переранжировщик в конвейере генерации с дополнением поиском (RAG). В репозитории Stanford ColBERT реализованы два решения. PLAID — это поисковая система, которая оценивает кандидатов по центроидам кластеров перед запуском полного MaxSim. ColBERTv2 (NAACL 2022) хранит каждый вектор токена как ID ближайшего центроида и остаток, квантованный до 1 или 2 бит на измерение; модель обучается на очищенных метках, полученных дистилляцией из кросс-энкодера. Сжатие остатков уменьшило объём памяти для позднего взаимодействия в 6–10 раз.