Turbovec: Thư viện tìm kiếm vector hiệu năng cao bằng Rust từ Google

Một tập dữ liệu gồm 10 triệu tài liệu chiếm 31 GB RAM ở định dạng float32. turbovec nén nó xuống còn 4 GB - và tìm kiếm nhanh hơn FAISS.

turbovec là một chỉ mục vector viết bằng Rust với các liên kết Python, được xây dựng dựa trên thuật toán TurboQuant của Google Research — một bộ lượng tử hóa không phụ thuộc vào dữ liệu với độ méo gần như tối ưu và không cần giai đoạn huấn luyện riêng biệt.

Bạn đang xây dựng RAG mà ở đó quyền riêng tư, bộ nhớ hoặc độ trễ là yếu tố quan trọng? Bạn đã đến đúng nơi rồi đấy.

Các vector và truy vấn là mảng 2-D float32 có hình dạng (n, dim) — các kiểu dữ liệu khác sẽ bị từ chối thay vì tự động chuyển đổi ngầm định, vì vậy hãy ép kiểu bằng np.asarray(x, dtype=np.float32) trước nếu cần.

Cần các ID ổn định có thể tồn tại sau khi xóa? Hãy sử dụng IdMapIndex:

Truy xuất lai (tìm kiếm có lọc)

Giới hạn kết quả trong một tập ứng viên được tạo ra bởi một hệ thống khác (SQL, BM25, ACL, cửa sổ thời gian, …):

Việc lọc diễn ra bên trong nhân SIMD với độ chi tiết khối 32-vector: các khối không có vị trí hợp lệ sẽ bị ngắt mạch trước khi thực hiện bất kỳ thao tác tra cứu LUT hoặc tính điểm nào, và các vị trí không hợp lệ riêng lẻ bên trong các khối đã tính điểm sẽ bị loại bỏ tại bước chèn vào heap. Do đó, các danh sách cho phép chọn lọc (chỉ một phần nhỏ của chỉ mục được cho phép) giúp tránh hầu hết chi phí SIMD thay vì phải trả phí rồi mới loại bỏ kết quả sau đó.

Độ dài đầu ra là min(k, n_allowed), trong đó n_allowed đếm các vector hợp lệ riêng biệt — khi số lượng vector được phép ít hơn k, bạn sẽ nhận được chính xác số lượng kết quả đó thay vì các kết quả dự phòng được đệm thêm.

Xem docs/api.md để biết tài liệu tham khảo đầy đủ.

Tích hợp khung làm việc (Framework)

Thay thế trực tiếp cho các kho lưu trữ vector/tài liệu tham chiếu có sẵn trong mỗi framework. Cùng bề mặt công khai, cùng ngữ nghĩa lưu trữ, cùng cách kết nối bộ truy xuất và đường ống — chỉ cần thay đổi import và giữ nguyên đường ống của bạn.

Rust

Đối với các ID bên ngoài ổn định có thể tồn tại sau khi xóa:

Độ hồi tưởng (Recall)

TurboQuant so với FAISS IndexPQ (LUT256, nbits=8) — đường cơ sở tại Mục 4.4 của bài báo. 100K vector, k=64. Số lượng bộ lượng tử hóa con của FAISS PQ được điều chỉnh để khớp với tốc độ bit của TurboQuant (m=d/4 ở 2-bit, m=d/2 ở 4-bit).

Các biểu đồ vẽ TurboQuant đã hiệu chuẩn (TQ+). Trên các tập OpenAI d=1536 và d=3072, TQ+ vượt qua FAISS ở chỉ số R@1 trên ba trong bốn ô (tăng 0,9–2,9 điểm; d=1536 4-bit thấp hơn 0,7), và cả hai đều đạt 1.0 tại k=8 (đã đạt ≥0,997 tại k≤4). GloVe d=200 là môi trường khó khăn hơn — ở số chiều thấp, giả định Beta tiệm cận trở nên lỏng lẻo hơn. TQ+ dẫn trước FAISS ở R@1 tại cả hai độ rộng bit (+1,9 ở 4-bit, +0,8 ở 2-bit), với FAISS giữ lợi thế nhỏ ở 2-bit từ k≈8. Các con số chưa hiệu chuẩn nằm trong các tệp JSON (tq_recalls).

Một lưu ý về các đường cơ sở. Chúng tôi so sánh với FAISS IndexPQ (LUT256, nbits=8, float32 LUT) vì đây là PQ cấp sản xuất mặc định mà hầu hết người dùng sẽ chọn. Đây là đường cơ sở mạnh hơn so với PQ sử dụng u8-LUT tùy chỉnh trong bài báo TurboQuant — FAISS sử dụng LUT có độ chính xác cao hơn khi tính điểm và k-means++ để huấn luyện bộ mã. Chúng tôi tái tạo các con số TurboQuant của bài báo trên OpenAI d=1536 / d=3072 và đạt được các con số tương tự như các triển khai tham chiếu cộng đồng khác trên các embedding số chiều thấp (xem turboquant-py tại d=384). Trên GloVe (d=200) — môi trường số chiều thấp nơi giả định Beta tiệm cận lỏng lẻo nhất — TurboQuant dẫn trước FAISS ở 4-bit nhưng lại thấp hơn ở 2-bit; hiệu chuẩn TQ+ đã bù đắp được sự thiếu hụt 2-bit tại R@1 (0,572 so với 0,564 của FAISS), trong khi FAISS giữ lợi thế nhỏ ở k sâu hơn.

Kết quả đầy đủ: d=1536 2-bit, d=1536 4-bit, d=3072 2-bit, d=3072 4-bit, GloVe 2-bit, GloVe 4-bit.

Nén

Tốc độ tìm kiếm

Tất cả các điểm chuẩn: 100K vector, 1K truy vấn, k=64, trung vị của 5 lần chạy.

ARM (GCP c4a-standard-8, Google Axion, 8 vCPU)

Trên ARM, TurboQuant vượt qua FAISS FastScan trong mọi cấu hình, trung bình 3,5× ở 4-bit (3,4–3,7× trên các ô — các nhân tích vô hướng SDOT/SMMLA tính toán trực tiếp trên bố cục vector-major) và 26% ở 2-bit (22–29%).

x86 (Intel Xeon Platinum 8481C / Sapphire Rapids, 8 vCPU)

Trên x86, TurboQuant thắng mọi cấu hình, trung bình 3,4× ở 4-bit (3,2–3,5× trên các ô — nhân tích vô hướng AVX-512 VNNI trên bố cục vector-major) và 20% ở 2-bit (5–32%), nơi quá trình quét LUT vpermb thực hiện vòng lặp tích lũy 2-bit ngắn.

Độ trễ chèn & xóa

Cùng tập dữ liệu như các ô tìm kiếm: 100K vector OpenAI, trung vị của 5 lần chạy, các vòng lặp tính thời gian bao gồm cả chi phí gọi Python mà người gọi thực sự phải trả cho mỗi thao tác. Chèn đo độ trễ add trên mỗi vector trên một chỉ mục đã được làm nóng, đã có dữ liệu (được xây dựng không tính thời gian) tại n=1 — một lần add đơn lẻ — và n=100 — một lô 100 vector, cho thấy việc xử lý theo lô giúp khấu hao chi phí mỗi lần gọi như thế nào — so với add vào FAISS IndexPQFastScan đã được huấn luyện và có dữ liệu (huấn luyện không tính thời gian). Một lần add đơn lẻ mất 6,3–19,7 µs tùy thuộc vào ô (nhanh hơn 7,6–13,9× so với một lần add của FAISS), và một lô 100 vector giúp TurboQuant khấu hao xuống còn 4,6–16,3 µs/vector (nhanh hơn 4,6–15,1× so với cùng lô đó vào FAISS). Xóa đo độ trễ remove-by-id trên mỗi thao tác tại n=1 (tốc độ ổn định trên mỗi thao tác qua 1000 lần xóa) và n=100 (100 lần xóa đầu tiên trên một chỉ mục mới): IdMapIndex.remove(id) — O(1) swap-and-pop cộng với việc quản lý id-map — đạt 0,44–1,22 µs và 0,59–1,37 µs mỗi thao tác trên các ô. Cột FAISS là cùng một thao tác hiển thị với người dùng, remove_ids trên IndexIDMap qua IndexPQFastScan, vốn đóng gói lại các mã đã lưu trữ trong mỗi lần gọi: 0,19–1,02 s cho mỗi lần xóa đơn lẻ tại 100K, với chi phí tăng gấp đôi cùng với kích thước mã — đó là lý do tại sao các biểu đồ xóa sử dụng trục log-scale. Các biểu đồ hiển thị các ô đơn luồng (RAYON_NUM_THREADS=1); các ô _mt cũng được đo lường và khớp tại n=1, vì một lần add đơn lẻ là tuần tự. Scripts: benchmarks/suite/.

ARM (GCP c4a-standard-8, Google Axion, 8 vCPU)

Kết quả đầy đủ: d=1536 2-bit insert, d=1536 4-bit insert, d=3072 2-bit insert, d=3072 4-bit insert, và các tệp speed_remove_* và _mt tương ứng.

x86 (Intel Xeon Platinum 8481C / Sapphire Rapids, 8 vCPU)

Kết quả đầy đủ: d=1536 2-bit insert, d=1536 4-bit insert, d=3072 2-bit insert, d=3072 4-bit insert, và các tệp speed_remove_* và _mt tương ứng.

Lưu & Tải

Cùng tập dữ liệu như các ô tìm kiếm: 100K vector OpenAI, trung vị của 5 lần chạy. TurboQuant tuần tự hóa thành một tệp.tv duy nhất với fsync + đổi tên nguyên tử; FAISS là write_index / read_index trên IndexPQFastScan có độ chính xác tương đương (số lượng bộ lượng tử hóa con khớp với tốc độ bit của TurboQuant, như trong các ô tìm kiếm). Lưu (warm) là thao tác ghi sau khi tìm kiếm đã chạy, vì vậy bộ nhớ đệm bố cục khối đã được điền. Tải → tìm kiếm đầu tiên mở một chỉ mục mới và tính thời gian cho truy vấn đầu tiên — tách biệt việc giải tuần tự hóa thuần túy (bộ nhớ đệm trang luôn nóng, vì vậy đây là công việc bố cục, không phải I/O lưu trữ lạnh) khỏi chi phí truy vấn đầu tiên. Vòng lặp khứ hồi (round-trip) liên kết chu kỳ kiểm tra/khôi phục mà một kho lưu trữ embedding thực sự phải trả — thay đổi 1K vector → lưu → mở lại → phục vụ truy vấn đầu tiên; FAISS không có phép đo tương đương cho đường dẫn này, vì vậy nó chỉ được hiển thị cho TurboQuant. Trên các tải trọng nhỏ hơn, vòng lặp khứ hồi có thể thấp hơn mức ghi sau khi thay đổi ("dirty") riêng biệt: cả hai được tính thời gian trong các bước bộ thử nghiệm riêng biệt, và ở kích thước tệp nhỏ, fsync độc lập trong bước ghi dirty chiếm ưu thế và làm tăng thời gian — một hiện tượng đo lường của bộ thử nghiệm, không phải là lợi thế đóng gói lại trong đường dẫn kết hợp. Các ô đơn luồng ghim RAYON_NUM_THREADS=1. Scripts: benchmarks/suite/.