ColBERT

ColBERT (Tương tác muộn theo ngữ cảnh trên BERT) là một model truy xuất neural biểu diễn truy vấn và từng tài liệu bằng một embedding cho mỗi token, rồi chấm điểm tài liệu bằng tương tác muộn: mỗi token truy vấn tìm token tài liệu tương tự nhất, sau đó cộng các độ tương tự cực đại này, gọi là MaxSim, thành điểm liên quan. Omar Khattab và Matei Zaharia tại Stanford University đã giới thiệu ColBERT trong bài báo SIGIR 2020 ColBERT: Tìm kiếm đoạn văn hiệu quả và chính xác thông qua tương tác muộn theo ngữ cảnh trên BERT. Embedding token của tài liệu được tính một lần, ngoại tuyến, vì vậy khi truy vấn, công việc chỉ là tính độ tương tự đơn giản thay vì chạy BERT đầy đủ cho từng tài liệu ứng viên.

Vị trí của tương tác muộn#

Các hệ thống truy xuất neural khác nhau ở thời điểm truy vấn và tài liệu được kết hợp. Bi-encoder nén mỗi văn bản thành một vector duy nhất, cấu hình tiêu chuẩn cho tìm kiếm vector: vector tài liệu được tính trước và tìm kiếm nhanh, nhưng một vector làm mờ khác biệt giữa từng thuật ngữ. Cross-encoder đưa truy vấn và tài liệu qua một transformer cùng lúc, để attention so sánh mọi cặp token. Chúng có độ chính xác cao nhất và phải chạy một lần cho mỗi cặp truy vấn-tài liệu, nên chỉ phù hợp để xếp hạng lại danh sách ứng viên ngắn.

ColBERT mã hóa truy vấn và tài liệu riêng biệt như bi-encoder, nhưng giữ một vector 128 chiều cho mỗi token và chỉ cho hai phía tương tác ở bước MaxSim cuối cùng. Truy vấn được đệm bằng token mask của BERT đến độ dài cố định 32, giúp model có thêm vị trí để mở rộng truy vấn bằng các thuật ngữ liên quan. Trên các benchmark tìm kiếm đoạn văn MS MARCO và TREC CAR, ColBERT đạt kết quả cạnh tranh với các bộ xếp hạng dựa trên BERT, đồng thời chạy nhanh hơn hai bậc độ lớn và dùng ít hơn bốn bậc độ lớn FLOPs cho mỗi truy vấn.

Tính điểm MaxSim#

Toán tử MaxSim chỉ cần vài dòng PyTorch. Ví dụ này áp dụng toán tử lên embedding token theo ngữ cảnh từ một encoder nhỏ thuộc họ 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}")

Đoạn văn về tiểu thuyết đạt 5.90 điểm và đoạn văn về cồn cát đạt 4.33 điểm. Cả hai đều chứa từ "dune", khớp mạnh ở mỗi đoạn, nhưng trong đoạn đầu tiên, token truy vấn "wrote" khớp với "published" với độ tương tự cosine là 0.72 và "novel" khớp với chính nó, trong khi ở đoạn thứ hai, kết quả khớp tốt nhất cho "wrote" chỉ đạt 0.29.

ColBERT và các thuật ngữ liên quan#

• Bản thân BERT: Encoder được tiền huấn luyện mà ColBERT tinh chỉnh cho tác vụ truy xuất, bổ sung một layer tuyến tính chiếu từng token thành 128 chiều.
• Late chunking: Cũng mã hóa văn bản trong ngữ cảnh trước khi chia nhỏ, sau đó lấy trung bình các token thành một vector cho mỗi chunk, trong khi ColBERT giữ lại mọi vector token.
• Tìm kiếm ngữ nghĩa: Tác vụ mà ColBERT thực hiện, xếp hạng đoạn văn theo ý nghĩa thay vì theo từ khóa chung.
• Tìm kiếm kết hợp: Kết hợp kết quả từ khóa và vector đơn để vừa khớp chính xác thuật ngữ vừa nắm bắt ý nghĩa. Cơ chế khớp ở cấp token của ColBERT thu được một phần khả năng khớp chính xác đó trong một model neural duy nhất.

Tương tác muộn cho hình ảnh#

Tương tác muộn cũng áp dụng được cho hình ảnh. ColPali (2024) nhúng ảnh chụp màn hình từng trang tài liệu thành lưới vector patch bằng model thị giác-ngôn ngữ và chấm điểm các vector này với token truy vấn bằng MaxSim, đưa tương tác muộn vào RAG đa phương thức trên PDF. Phát hiện từ vựng mở giữ đặc trưng hình ảnh ở mức chi tiết vì lý do tương tự: Ultralytics YOLOE lưu embedding văn bản cho từng lớp được nhắc đến, chẳng hạn "double-decker bus", rồi so sánh embedding đó với đặc trưng vùng tại mọi vị trí trong feature map, nhờ vậy một prompt ngắn có thể khớp với một vật thể nhỏ thay vì dùng một vector duy nhất tóm tắt toàn bộ hình ảnh.

ColBERTv2 và PLAID#

Lưu một vector cho mỗi token khiến chỉ mục ColBERT lớn hơn khoảng một bậc độ lớn so với chỉ mục vector đơn, một chi phí đáng kể khi ColBERT được dùng làm bộ truy xuất giai đoạn đầu hoặc bộ xếp hạng lại trong pipeline sinh tăng cường truy xuất (RAG). Kho lưu trữ Stanford ColBERT triển khai hai giải pháp. PLAID là công cụ tìm kiếm chấm điểm ứng viên dựa trên tâm cụm trước khi chạy MaxSim đầy đủ. ColBERTv2 (NAACL 2022) lưu mỗi vector token dưới dạng ID của tâm cụm gần nhất cộng với phần dư được lượng tử hóa xuống 1 hoặc 2 bit mỗi chiều, đồng thời huấn luyện bằng giám sát đã khử nhiễu được chưng cất từ cross-encoder; kỹ thuật nén phần dư giúp giảm dung lượng bộ nhớ của tương tác muộn từ 6 đến 10 lần.