Что такое ColBERT и позднее взаимодействие, и почему они важны для поиска?

Jina ColBERT v2: Многоязычный ретривер с поздним взаимодействием для эмбеддингов и переранжирования

Jina ColBERT v2 поддерживает 89 языков с превосходной производительностью поиска, настраиваемыми выходными размерностями и длиной токенов 8192.

В прошлую пятницу выпуск модели ColBERT от Jina AI на Hugging Face вызвал значительный интерес в AI-сообществе, особенно в Twitter/X. В то время как многие знакомы с революционной моделью BERT, ажиотаж вокруг ColBERT заставил некоторых задаться вопросом: Чем ColBERT выделяется среди множества технологий информационного поиска? Почему AI-сообщество в восторге от ColBERT с длиной 8192? Эта статья углубляется в тонкости ColBERT и ColBERTv2, освещая их дизайн, улучшения и удивительную эффективность позднего взаимодействия ColBERT.

Reranker API

Максимизируйте релевантность поиска и точность RAG с легкостью

tagЧто такое ColBERT?

Название "ColBERT" расшифровывается как Contextualized Late Interaction over BERT (Контекстуализированное позднее взаимодействие поверх BERT), модель, разработанная в Стэнфордском университете, которая использует глубокое понимание языка BERT, представляя новый механизм взаимодействия. Этот механизм, известный как позднее взаимодействие, позволяет осуществлять эффективный и точный поиск путем раздельной обработки запросов и документов до финальных этапов процесса поиска. Конкретно существует две версии модели:

• ColBERT: Первоначальная модель была детищем Omar Khattab и Matei Zaharia, представляющая новый подход к информационному поиску через статью "ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT". Их работа была опубликована на SIGIR 2020.

ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT

Recent progress in Natural Language Understanding (NLU) is driving fast-paced advances in Information Retrieval (IR), largely owed to fine-tuning deep language models (LMs) for document ranking. While remarkably effective, the ranking models based on these LMs increase computational cost by orders of magnitude over prior approaches, particularly as they must feed each query-document pair through a massive neural network to compute a single relevance score. To tackle this, we present ColBERT, a novel ranking model that adapts deep LMs (in particular, BERT) for efficient retrieval. ColBERT introduces a late interaction architecture that independently encodes the query and the document using BERT and then employs a cheap yet powerful interaction step that models their fine-grained similarity. By delaying and yet retaining this fine-granular interaction, ColBERT can leverage the expressiveness of deep LMs while simultaneously gaining the ability to pre-compute document representations offline, considerably speeding up query processing. Beyond reducing the cost of re-ranking the documents retrieved by a traditional model, ColBERT's pruning-friendly interaction mechanism enables leveraging vector-similarity indexes for end-to-end retrieval directly from a large document collection. We extensively evaluate ColBERT using two recent passage search datasets. Results show that ColBERT's effectiveness is competitive with existing BERT-based models (and outperforms every non-BERT baseline), while executing two orders-of-magnitude faster and requiring four orders-of-magnitude fewer FLOPs per query.

arXiv.orgOmar Khattab

• ColBERTv2: Основываясь на фундаментальной работе, Omar Khattab продолжил свои исследования, сотрудничая с Barlas Oguz, Matei Zaharia и Michael S. Bernstein, чтобы представить "ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction" на SIGIR 2021. Эта следующая итерация ColBERT решала предыдущие ограничения и вводила ключевые улучшения, такие как очищенное обучение и остаточное сжатие, повышая как эффективность поиска модели, так и её эффективность хранения.

ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction

Neural information retrieval (IR) has greatly advanced search and other knowledge-intensive language tasks. While many neural IR methods encode queries and documents into single-vector representations, late interaction models produce multi-vector representations at the granularity of each token and decompose relevance modeling into scalable token-level computations. This decomposition has been shown to make late interaction more effective, but it inflates the space footprint of these models by an order of magnitude. In this work, we introduce ColBERTv2, a retriever that couples an aggressive residual compression mechanism with a denoised supervision strategy to simultaneously improve the quality and space footprint of late interaction. We evaluate ColBERTv2 across a wide range of benchmarks, establishing state-of-the-art quality within and outside the training domain while reducing the space footprint of late interaction models by 6--10$\times$.

arXiv.orgKeshav Santhanam

tagПонимание дизайна ColBERT

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

tagЧто такое позднее взаимодействие в ColBERT?

"Взаимодействие" относится к процессу оценки релевантности между запросом и документом путем сравнения их представлений.

"Позднее взаимодействие" — это суть ColBERT. Термин происходит от архитектуры модели и стратегии обработки, где взаимодействие между представлениями запроса и документа происходит поздно в процессе, после того как оба были независимо закодированы. Это контрастирует с моделями "раннего взаимодействия", где эмбеддинги запроса и документа взаимодействуют на более ранних этапах, обычно до или во время их кодирования моделью.

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

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

tagОтсутствие взаимодействия: косинусное сходство эмбеддингов документа и запроса

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

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

tagКодировщики запросов и документов в ColBERT

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

tagКодировщик запросов ColBERT

Для запроса $Q$ с токенами ${q_1, q_2, ..., q_l}$ процесс начинается с токенизации $Q$ в токены WordPiece на основе BERT и добавления специального токена [Q]. Этот токен [Q], расположенный сразу после токена [CLS] BERT, сигнализирует о начале запроса.

Если запрос короче предопределенного количества токенов $N_q$, он дополняется токенами [mask] до $N_q$; в противном случае он усекается до первых $N_q$ токенов. Дополненная последовательность затем проходит через BERT, за которым следует CNN (Сверточная Нейронная Сеть) и нормализация. Результатом является набор векторов эмбеддингов, обозначенный как $\mathbf{E}_q$ ниже:
$$\mathbf{E}_q := \mathrm{Normalize}\left(\mathrm{BERT}\left(\mathtt{[Q]},q_0,q_1,\ldots,q_l\mathtt{[mask]},\mathtt{[mask]},\ldots,\mathtt{[mask]}\right)\right)$$

tagКодировщик документов ColBERT

Аналогично, для документа $D$ с токенами ${d_1, d_2, ..., d_n}$ токен [D] добавляется в начало для обозначения начала документа. Эта последовательность, не требующая дополнения, проходит тот же процесс, что приводит к набору векторов эмбеддингов, обозначенному как $\mathbf{E}_d$ ниже:
$$\mathbf{E}_d := \mathrm{Filter}\left(\mathrm{Normalize}\left(\mathrm{BERT}\left(\mathtt{[D]},d_0,d_1,...,d_n\right)\right)\right)$$

Использование токенов [mask] для дополнения запросов (названное в статье расширением запроса) обеспечивает единообразную длину всех запросов, облегчая пакетную обработку. Токены [Q] и [D] явно отмечают начало запросов и документов соответственно, помогая модели различать эти два типа входных данных.

tagСравнение ColBERT с кросс-энкодерами

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

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

tagПоиск топ-K документов с помощью ColBERT

После получения эмбеддингов для запроса и документов, поиск наиболее релевантных топ-K документов становится прямолинейным (но не таким простым, как вычисление косинуса двух векторов).

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

import torch

def compute_relevance_scores(query_embeddings, document_embeddings, k):
    """
    Compute relevance scores for top-k documents given a query.
    
    :param query_embeddings: Tensor representing the query embeddings, shape: [num_query_terms, embedding_dim]
    :param document_embeddings: Tensor representing embeddings for k documents, shape: [k, max_doc_length, embedding_dim]
    :param k: Number of top documents to re-rank
    :return: Sorted document indices based on their relevance scores
    """
    
    # Ensure document_embeddings is a 3D tensor: [k, max_doc_length, embedding_dim]
    # Pad the k documents to their maximum length for batch operations
    # Note: Assuming document_embeddings is already padded and moved to GPU
    
    # Compute batch dot-product of Eq (query embeddings) and D (document embeddings)
    # Resulting shape: [k, num_query_terms, max_doc_length]
    scores = torch.matmul(query_embeddings.unsqueeze(0), document_embeddings.transpose(1, 2))
    
    # Apply max-pooling across document terms (dim=2) to find the max similarity per query term
    # Shape after max-pool: [k, num_query_terms]
    max_scores_per_query_term = scores.max(dim=2).values
    
    # Sum the scores across query terms to get the total score for each document
    # Shape after sum: [k]
    total_scores = max_scores_per_query_term.sum(dim=1)
    
    # Sort the documents based on their total scores
    sorted_indices = total_scores.argsort(descending=True)
    
    return sorted_indices
Копировать

Обратите внимание, что эта процедура используется как при обучении, так и при переранжировании во время вывода. Модель ColBERT обучается с использованием функции потерь попарного ранжирования, где обучающие данные состоят из троек $(q, d^+, d^-)$, где $q$ представляет запрос, $d^+$ - релевантный (положительный) документ для запроса, а $d^-$ - нерелевантный (отрицательный) документ. Модель стремится научиться создавать такие представления, чтобы оценка сходства между $q$ и $d^+$ была выше, чем оценка между q и $d^-$.

Цель обучения математически может быть представлена как минимизация следующей функции потерь: $$\mathrm{Loss} = \max(0, 1 - S(q, d^+) + S(q, d^-))$$

, где $S(q, d)$ обозначает оценку сходства, вычисленную ColBERT между запросом $q$ и документом $d$. Эта оценка получается путем агрегации максимальных оценок сходства наилучших совпадающих эмбеддингов между запросом и документом, следуя паттерну позднего взаимодействия, описанному в архитектуре модели. Этот подход гарантирует, что модель обучается различать релевантные и нерелевантные документы для данного запроса, поощряя больший разрыв в оценках сходства для положительных и отрицательных пар документов.

tagОчищенное обучение в ColBERTv2

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

1. Начальное обучение: Использование официальных троек из набора данных MS MARCO, состоящих из запроса, релевантного документа и нерелевантного документа.
2. Индексация и поиск: Применение сжатия ColBERTv2 для индексации обучающих пассажей с последующим извлечением top-k пассажей для каждого запроса.
3. Переранжирование кросс-энкодером: Улучшение выбора пассажей путем переранжирования с помощью кросс-энкодера MiniLM, дистиллируя его оценки в ColBERTv2.
4. Формирование обучающих кортежей: Генерация w-way кортежей для обучения, включающих как высоко-, так и низкоранжированные пассажи для создания сложных примеров.
5. Итеративное уточнение: Повторение процесса для постоянного улучшения выбора сложных негативных примеров, тем самым повышая производительность модели.

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

tagГиперпараметры ColBERT

Гиперпараметры ColBERT приведены ниже:

tagСтратегия индексирования ColBERT

В отличие от подходов, основанных на представлении, которые кодируют каждый документ в один вектор эмбеддинга, ColBERT кодирует документы (и запросы) в наборы эмбеддингов, где каждый токен в документе имеет свой собственный эмбеддинг. Этот подход по своей сути означает, что для более длинных документов будет храниться больше эмбеддингов, что является слабым местом оригинального ColBERT, которое позже было решено в ColBERTv2.

Ключ к эффективному управлению этим заключается в использовании ColBERT векторной базы данных (например, FAISS) для индексирования и поиска, и его детальном процессе индексирования, который разработан для эффективной обработки больших объемов данных. В оригинальной статье о ColBERT упоминается несколько стратегий для повышения эффективности индексирования и поиска, включая:

• Офлайн индексирование: Представления документов вычисляются офлайн, позволяя предварительно вычислять и хранить эмбеддинги документов. Этот процесс использует пакетную обработку и ускорение GPU для эффективной обработки больших коллекций документов.
• Хранение эмбеддингов: Эмбеддинги документов могут храниться с использованием 32-битных или 16-битных значений для каждого измерения, предлагая компромисс между точностью и требованиями к хранению. Эта гибкость позволяет ColBERT поддерживать баланс между эффективностью (с точки зрения производительности поиска) и эффективностью (с точки зрения затрат на хранение и вычисления).

Введение остаточного сжатия в ColBERTv2, которое является новым подходом, отсутствующим в оригинальном ColBERT, играет ключевую роль в уменьшении пространственного следа модели в 6-10 раз при сохранении качества. Эта техника дополнительно сжимает эмбеддинги, эффективно захватывая и сохраняя только различия от набора фиксированных опорных центроидов.

tagЭффективность и производительность ColBERT

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

1. Значительный прирост эффективности: ColBERT достигает снижения вычислительных затрат (FLOP) и задержки на порядки по сравнению с традиционными моделями ранжирования на основе BERT. В частности, для данного размера модели (например, 12-слойный "базовый" трансформер-энкодер), ColBERT не только соответствует, но в некоторых случаях превосходит эффективность моделей на основе BERT при значительно меньших вычислительных требованиях. Например, при глубине повторного ранжирования k=10, BERT требует почти в 180 раз больше FLOP, чем ColBERT; этот разрыв увеличивается с ростом k, достигая 13900× при k=1000 и даже 23000× при k=2000.
2. Улучшенные Recall и MRR@10 в end-to-end поиске: Вопреки начальной интуиции о том, что более глубокое взаимодействие между представлениями запроса и документа (как в моделях раннего взаимодействия) будет необходимо для высокой эффективности поиска, end-to-end схема поиска ColBERT демонстрирует превосходную эффективность. Например, его Recall@50 превосходит официальный BM25 Recall@1000 и почти все другие модели Recall@200, подчеркивая замечательную способность модели извлекать релевантные документы из обширной коллекции без прямого сравнения каждой пары запрос-документ.
3. Практичность для реальных приложений: Экспериментальные результаты подчеркивают практическую применимость ColBERT для реальных сценариев. Его пропускная способность индексирования и эффективность использования памяти делают его подходящим для индексирования больших коллекций документов, таких как MS MARCO, в течение нескольких часов, сохраняя высокую эффективность при управляемом использовании пространства. Эти качества подчеркивают пригодность ColBERT для развертывания в производственных средах, где важны как производительность, так и вычислительная эффективность.
4. Масштабируемость с размером коллекции документов: Возможно, самым удивительным выводом является масштабируемость и эффективность ColBERT в обработке крупномасштабных коллекций документов. Архитектура позволяет предварительно вычислять эмбеддинги документов и использует эффективную пакетную обработку для взаимодействия запрос-документ, позволяя системе эффективно масштабироваться с размером коллекции документов. Эта масштабируемость противоречит интуиции при учете сложности и глубины понимания, необходимых для эффективного поиска документов, демонстрируя инновационный подход ColBERT к балансированию вычислительной эффективности с эффективностью поиска.

tagИспользование jina-colbert-v1-en: модель ColBERTv2 с длиной контекста 8192

Jina-ColBERT разработан для быстрого и точного поиска, поддерживая контексты длиной до 8192 токенов, используя преимущества JinaBERT, который позволяет обрабатывать более длинные последовательности благодаря улучшениям в архитектуре.

jinaai/jina-colbert-v1-en · Hugging Face

We're on a journey to advance and democratize artificial intelligence through open source and open science.

tagУлучшения Jina по сравнению с оригинальным ColBERT

Основным улучшением Jina-ColBERT является его основа, jina-bert-v2-base-en, которая позволяет обрабатывать значительно более длинные контексты (до 8192 токенов) по сравнению с оригинальным ColBERT, использующим bert-base-uncased. Эта возможность крайне важна для обработки документов с обширным содержанием, обеспечивая более детальные и контекстуальные результаты поиска.

tagСравнение производительности jina-colbert-v1-en с ColBERTv2

Мы оценили jina-colbert-v1-en на датасетах BEIR и новом бенчмарке LoCo, который ориентирован на длинный контекст, протестировав его против оригинальной реализации ColBERTv2 и основанных на отсутствии взаимодействия

Эта таблица демонстрирует превосходную производительность jina-colbert-v1-en，особенно в сценариях，требующих большей длины контекста по сравнению с оригинальным ColBERTv2. Стоит отметить，что jina-embeddings-v2-base-en использует больше обучающих данных，в то время как jina-colbert-v1-en использует только MSMARCO，что может объяснять хорошую производительность jina-embeddings-v2-base-en в некоторых задачах.

tagПример использования jina-colbert-v1-en

Этот фрагмент кода описывает процесс индексации с помощью Jina-ColBERT，демонстрируя поддержку длинных документов.

from colbert import Indexer
from colbert.infra import Run, RunConfig, ColBERTConfig

n_gpu: int = 1  # Set your number of available GPUs
experiment: str = ""  # Name of the folder where the logs and created indices will be stored
index_name: str = ""  # The name of your index, i.e. the name of your vector database

if __name__ == "__main__":
    with Run().context(RunConfig(nranks=n_gpu, experiment=experiment)):
        config = ColBERTConfig(
          doc_maxlen=8192  # Our model supports 8k context length for indexing long documents
        )
        indexer = Indexer(
          checkpoint="jinaai/jina-colbert-v1-en",
          config=config,
        )
        documents = [
          "ColBERT is an efficient and effective passage retrieval model.",
          "Jina-ColBERT is a ColBERT-style model but based on JinaBERT so it can support both 8k context length.",
          "JinaBERT is a BERT architecture that supports the symmetric bidirectional variant of ALiBi to allow longer sequence length.",
          "Jina-ColBERT model is trained on MSMARCO passage ranking dataset, following a very similar training procedure with ColBERTv2.",
          "Jina-ColBERT achieves the competitive retrieval performance with ColBERTv2.",
          "Jina is an easier way to build neural search systems.",
          "You can use Jina-ColBERT to build neural search systems with ease.",
          # Add more documents here to ensure the clustering work correctly
        ]
        indexer.index(name=index_name, collection=documents)
Копировать

tagИспользование jina-colbert-v1-en в RAGatouille

RAGatouille — это новая библиотека Python，которая упрощает использование продвинутых методов поиска в RAG-пайплайнах. Она разработана с учетом модульности и простоты интеграции，позволяя пользователям легко применять передовые исследования. Основная цель RAGatouille — упростить применение сложных моделей，таких как ColBERT，в RAG-пайплайнах，делая их доступными для разработчиков без необходимости глубокого понимания базовых исследований. Благодаря Benjamin Clavié，теперь вы можете легко использовать jina-colbert-v1-en:

from ragatouille import RAGPretrainedModel

# Get your model & collection of big documents ready
RAG = RAGPretrainedModel.from_pretrained("jinaai/jina-colbert-v1-en")
my_documents = [
    "very long document1",
    "very long document2",
    # ... more documents
]

# And create an index with them at full length!
RAG.index(collection=my_documents,
          index_name="the_biggest_index",
          max_document_length=8190,)

# or encode them in-memory with no truncation, up to your model's max length
RAG.encode(my_documents)
Копировать

Для получения более подробной информации о Jina-ColBERT вы можете посетить страницу на Hugging Face.

tagЗаключение

ColBERT представляет собой значительный шаг вперед в области информационного поиска. Благодаря возможности работы с более длинным контекстом в Jina-ColBERT и сохранению совместимости с подходом ColBERT к позднему взаимодействию，он предлагает мощную альтернативу для разработчиков，стремящихся реализовать современную функциональность поиска.

В сочетании с библиотекой RAGatouille，которая упрощает интеграцию сложных моделей поиска в RAG-пайплайны，разработчики теперь могут легко использовать возможности продвинутого поиска，оптимизируя свои рабочие процессы и улучшая свои приложения. Синергия между Jina-ColBERT и RAGatouille демонстрирует замечательный шаг в направлении обеспечения доступности и эффективности продвинутых моделей поиска на базе ИИ для практического использования.