ColBERT

ColBERT (Interação tardia contextualizada sobre BERT) é um modelo neural de recuperação que representa uma consulta e cada documento com um embedding por token e pontua um documento por meio de interação tardia: cada token da consulta encontra o token mais semelhante do documento, e essas similaridades máximas, chamadas MaxSim, são somadas para formar a pontuação de relevância. Omar Khattab e Matei Zaharia, da Universidade Stanford, o apresentaram no artigo da SIGIR 2020 ColBERT: Busca de passagens eficiente e eficaz por meio de interação tardia contextualizada sobre BERT. Os embeddings dos tokens do documento são calculados uma vez, offline; assim, o processamento no momento da consulta consiste em um cálculo de similaridade simples, em vez de uma passagem completa pelo BERT para cada documento candidato.

Onde a interação tardia se encaixa#

Os recuperadores neurais diferem em quando a consulta e o documento se encontram. Bi-encoders comprimem cada texto em um único vetor, a configuração padrão para busca vetorial: os vetores dos documentos são pré-calculados e a busca é rápida, mas um único vetor confunde as distinções entre termos individuais. Cross-encoders processam juntos a consulta e o documento por meio de um Transformer, de modo que a atenção compara cada par de tokens. São os mais precisos e precisam ser executados uma vez para cada par consulta-documento, o que os limita à reordenação de listas curtas de candidatos.

O ColBERT codifica a consulta e o documento separadamente, como faz um bi-encoder, mas mantém um vetor de 128 dimensões por token e permite que os dois lados interajam somente na etapa final MaxSim. As consultas são preenchidas com o token de máscara do BERT até um comprimento fixo de 32, o que dá ao modelo posições adicionais para expandir a consulta com termos relacionados. Nos benchmarks de busca de passagens MS MARCO e TREC CAR, o ColBERT foi competitivo com os classificadores baseados em BERT, executando duas ordens de grandeza mais rápido e usando quatro ordens de grandeza menos FLOPs por consulta.

Calculando uma pontuação MaxSim#

O operador MaxSim pode ser implementado em poucas linhas de PyTorch. Este exemplo o aplica a embeddings contextuais de tokens de um pequeno codificador da família 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}")

A passagem sobre o romance recebe a pontuação 5,90, e a passagem sobre a duna de areia, 4,33. Ambas contêm “duna”, que apresenta uma correspondência forte em cada uma; porém, na primeira passagem, o token “escreveu” da consulta encontra “publicou” com uma similaridade do cosseno de 0,72, e “romance” encontra a si próprio, enquanto, na segunda passagem, a melhor correspondência para “escreveu” chega apenas a 0,29.

ColBERT em comparação com termos relacionados#

• O próprio BERT: o codificador pré-treinado que o ColBERT ajusta para recuperação, adicionando uma camada linear que projeta cada token em 128 dimensões.
• Fragmentação tardia: também codifica o texto em contexto antes de dividi-lo e, em seguida, calcula a média dos tokens em um vetor por fragmento, enquanto o ColBERT mantém cada vetor de token.
• Busca semântica: a tarefa realizada pelo ColBERT, que classifica passagens por significado, em vez de palavras-chave compartilhadas.
• Busca híbrida: combina resultados de palavras-chave e de vetores únicos para reunir correspondências exatas de termos e significado. A correspondência em nível de token do ColBERT captura parte dessa exatidão dentro de um único modelo neural.

Interação tardia para imagens#

A interação tardia também funciona com imagens. O ColPali (2024) incorpora a captura de tela de cada página de documento como uma grade de vetores de recortes usando um modelo de visão e linguagem e os pontua em relação aos tokens da consulta com MaxSim, levando a interação tardia à RAG multimodal em PDFs. A detecção de vocabulário aberto mantém as características da imagem em granularidade fina por um motivo relacionado: o Ultralytics YOLOE armazena um embedding de texto para cada classe solicitada, como “ônibus de dois andares”, e o compara com as características das regiões em cada posição dos seus mapas de características, de modo que uma solicitação curta pode corresponder a um único objeto pequeno, em vez de a um único vetor que resume a imagem inteira.

ColBERTv2 e PLAID#

Armazenar um vetor para cada token faz com que um índice ColBERT seja aproximadamente uma ordem de grandeza maior que um índice de vetor único, um custo relevante quando o ColBERT funciona como recuperador ou reclassificador de primeira etapa em um pipeline de geração aumentada por recuperação (RAG). O repositório Stanford ColBERT implementa duas soluções. O PLAID é um mecanismo de busca que pontua candidatos em relação a centroides de agrupamento antes de executar o MaxSim completo. O ColBERTv2 (NAACL 2022) armazena cada vetor de token como o ID do centroide mais próximo, além de um resíduo quantizado a 1 ou 2 bits por dimensão, e treina com supervisão sem ruído destilada de um cross-encoder; sua compressão residual reduziu o espaço ocupado pela interação tardia de 6 a 10 vezes.