TurboQuant: Kỹ thuật lượng tử hóa tối ưu cho mô hình AI và Tìm kiếm Vector

Trong bối cảnh các mô hình ngôn ngữ lớn (LLM) và hệ thống tìm kiếm vector ngày càng phổ biến, việc tối ưu hóa bộ nhớ thông qua lượng tử hóa (quantization) trở nên vô cùng quan trọng. TurboQuant là một kỹ thuật mới nổi bật, giải quyết bài toán nén vector theo cách tiếp cận từ những nguyên lý cơ bản nhất của toán học, mang lại hiệu suất vượt trội so với các phương pháp truyền thống.

Vấn đề của lượng tử hóa truyền thống

Lượng tử hóa vector (Vector Quantization) là quá trình giảm độ chính xác của các số thực trong một vector để tiết kiệm dung lượng lưu trữ, thường diễn ra ở các tầng nhúng (embeddings) hoặc bộ nhớ đệm KV (KV cache) của LLM. Các phương pháp lượng tử hóa ngây thơ thường sử dụng một lưới cố định để "bắt" các giá trị về mức gần nhất.

Tuy nhiên, thực tế cho thấy các vector nhúng thường chứa các "giá trị ngoại lệ" (outliers) — một vài thành phần có độ lớn lớn hơn nhiều so với phần còn lại. Nếu sử dụng lưới cố định, các phương pháp này sẽ hoặc là cắt cụt các giá trị ngoại lệ, hoặc là lãng phí độ phân giải cho các giá trị nhỏ. Các giải pháp sản xuất hiện tại như GPTQ hay AWQ thường phải tính toán các thông số điều chỉnh theo từng khối nhỏ và lưu trữ chúng dưới dạng siêu dữ liệu (metadata). Dù hiệu quả, nhưng cách này làm phát sinh thêm chi phí lưu trữ, khiến một quy trình "3-bit" thực tế lại tốn tới 4-5 bit cho mỗi giá trị.

Thủ thuật Xoay (The Rotation Trick)

TurboQuant giải quyết vấn đề này bằng một cách tiếp cận thông minh: xoay ngẫu nhiên. Bằng cách nhân vector đầu vào với một ma trận xoay trực giao ngẫu nhiên $\boldsymbol{\Pi}$, thuật toán phân phối lại độ lớn của vector đó khắp các chiều không gian.

Trong không gian nhiều chiều (high-dimensional space), phép xoay này làm cho bất kỳ vector đầu nào cũng trông giống nhau về mặt thống kê. Cụ thể, theo Định lý giới hạn trung tâm (CLT) và sự tập trung độ đo (measure concentration), các tọa độ của một vector được xoay ngẫu nhiên sẽ tuân theo phân phối xác suất cố định (tiệm cận phân phối Gauss). Điều này có nghĩa là "đỉnh nhọn" của giá trị ngoại lệ sẽ bị san phẳng đều đặn trên toàn bộ các chiều của vector.

Bảng mã phổ quát (Universal Codebook)

Nhờ vào thủ thuật xoay, TurboQuant biến bài toán lượng tử hóa vector phức tạp thành bài toán lượng tử hóa vô hướng đơn giản. Vì mọi tọa độ sau khi xoay đều có cùng một phân phối xác suất đã biết, chúng ta chỉ cần xây dựng một bảng mã (codebook) tối ưu duy nhất một lần duy nhất.

Để xây dựng bảng mã này, TurboQuant sử dụng thuật toán Lloyd-Max cổ điển. Thuật toán này tìm ra các điểm đại diện (centroids) sao cho việc làm tròn các mẫu về điểm gần nhất gây ra sai số bình phương trung bình (MSE) thấp nhất. Do bảng mã được tính toán sẵn và không phụ thuộc vào dữ liệu đầu vào cụ thể, quá trình mã hóa chỉ đơn giản là tra cứu bảng (lookup), giúp loại bỏ hoàn toàn chi phí siêu dữ liệu theo từng khối.

Khắc phục thiên lệch Tích vô hướng

Mặc dù tối ưu hóa MSE tốt, nhưng các bộ lượng tử hóa truyền thống thường tạo ra sự thu hẹp (shrinkage) đối với tích vô hướng (inner product) — một phép toán quan trọng trong cơ chế Attention của các mô hình Transformer. Vector được tái tạo thường có độ lớn nhỏ hơn vector gốc, dẫn đến việc đánh giá điểm số Attention bị thấp hơn thực tế (thiên lệch hệ thống).

TurboQuant giải quyết vấn đề này bằng cách kết hợp kỹ thuật QJL (Quantized Johnson-Lindenstrauss). Một phần ngân sách bit được dành riêng để mã hóa phần dư và loại bỏ thiên lệch này. Bộ giải mã sẽ nhân với một hệ số Calibration định sẵn để đảm bảo kỳ vọng của tích vô hướng bằng với giá trị thực, tạo ra một bộ ước lượng không thiên lệch (unbiased estimator).

Hiệu năng thực tế và Tác động

Khi áp dụng vào thực tế, TurboQuant mang lại những kết quả ấn tượng:

• Nén bộ nhớ đệm KV: Trên bài toán Needle-in-a-Haystack với Llama-3.1-8B, TurboQuant đạt điểm số nhớ hoàn hảo (0.997) ở mức nén 4x, tương đương với độ chính xác của bản đầy đủ (FP16).
• Tốc độ: Trong các bài toán tìm kiếm láng giềng gần nhất (nearest neighbor search), TurboQuant nhanh hơn các phương pháp như Product Quantization hay RabitQ từ 4 đến 6 cấp độ (magnitudes) nhờ việc loại bỏ quá trình học codebook từ dữ liệu và không cần tính toán scale theo khối tại thời điểm index.
• Tối ưu lý thuyết: Sai số tái tạo của TurboQuant nằm trong hằng số nhỏ (~1.45x ở mức 1-bit, ~2.72x tiệm cận) so với giới hạn dưới của Shannon, chứng minh hiệu quả gần như tối ưu về mặt lý thuyết thông tin.

TurboQuant đại diện cho một bước tiến quan trọng trong việc tối ưu hóa cơ sở hạ tầng AI, cho phép chạy các mô hình lớn với chi phí bộ nhớ thấp hơn mà không đánh đổi quá nhiều về chất lượng suy luận.