Bài 7: Nén Vector PQ (Product Quantization) & Tìm Kiếm Trên Dữ Liệu Nén

Một véc-tơ đơn lẻ của mô hình text-embedding-3-small chiếm 1536 số thực Float32, tức tương đương với 6KB bộ nhớ. Với 100 triệu vector, dung lượng RAM tối thiểu cần dùng là 600GB - một chi phí hạ tầng máy chủ vô cùng đắt đỏ. Để mở rộng hệ thống lên hàng tỷ véc-tơ, việc áp dụng các kỹ thuật nén véc-tơ là vô cùng cấp thiết.

Bài học này sẽ hướng dẫn bạn chi tiết thuật toán nén vector phổ biến hàng đầu hiện nay: Product Quantization (Lượng tử hóa sản phẩm - PQ), nguyên lý xây dựng Codebook để mã hóa vector thô thành một chuỗi byte nhỏ gọn, và giải thuật tìm kiếm bất đối xứng (ADC) trực tiếp trên dữ liệu đã nén.

7.1 Tại sao cần nén Vector & Nguyên lý lượng tử hóa

Lượng tử hóa (Quantization) là quá trình ánh xạ một tập hợp lớn các giá trị liên tục thành một tập hợp nhỏ các giá trị rời rạc.

Thay vì lưu trữ giá trị chính xác tuyệt đối của từng số thực Float32 (tốn 32 bits), chúng ta biểu diễn dữ liệu bằng các chỉ số đại diện ngắn hơn (như số nguyên Int8 tốn 8 bits). Sự đánh đổi ở đây là sai số lượng tử hóa (quantization noise), tức là mất đi một chút độ mịn của tọa độ véc-tơ ban đầu nhưng đổi lại dung lượng lưu trữ giảm đi đáng kể.

7.2 Thuật toán lượng tử hóa sản phẩm (Product Quantization - PQ)

Nếu áp dụng Scalar Quantization (SQ) thông thường trên từng số đơn lẻ, tỷ lệ nén tối đa chỉ khoảng 4 lần (từ 32-bit xuống 8-bit). Product Quantization (PQ) giải quyết bài toán này bằng cách gom nhóm các chiều lại và nén chung.

Quy trình PQ chia véc-tơ thành các phần nhỏ hơn:

1. Chia một vector $\mathbf{x}$ kích thước $d$ chiều thành $m$ phân không gian con (subspaces) độc lập, mỗi subspace chứa một vector con (sub-vector) kích thước $d^* = d / m$ chiều:
\[\mathbf{x} = [\mathbf{x}^1, \mathbf{x}^2, \dots, \mathbf{x}^m] \in \mathbb{R}^d\]
2. Chạy giải thuật K-Means riêng biệt trên từng phân không gian con này để gom các dữ liệu mẫu thành $k^*$ cụm. Mỗi cụm được đại diện bởi một centroid con.

// Chia nhỏ vector 128 chiều thành 16 sub-vectors 8 chiều
function splitSubspaces(vector, m = 16) {
  const subDim = vector.length / m;
  const subVectors = [];
  
  for (let i = 0; i < m; i++) {
    const start = i * subDim;
    subVectors.push(vector.slice(start, start + subDim));
  }
  return subVectors;
}

7.3 Xây dựng Codebook và Mã hóa vector

Tập hợp các centroids của tất cả $m$ phân không gian con tạo thành bảng mã gọi là Codebook.

Khi mã hóa (quantize) một vector mới:

1. Ta chia nó thành $m$ sub-vectors.
2. Trong mỗi sub-space thứ $i$, tìm centroid gần nhất trong Codebook với sub-vector tương ứng.
3. Lưu lại chỉ số (index) của centroid đó. Chỉ số này thường nằm từ $0$ đến $255$ nên chỉ tốn đúng 1 byte (Uint8) để lưu trữ.

Như vậy, một vector Float32 ban đầu đã được chuyển đổi thành một mảng gồm $m$ chỉ số Uint8. Kích thước giảm từ $d \times 4$ bytes xuống còn đúng $m$ bytes.

// Mã hóa vector thành chuỗi byte codes
function encodeVector(vector, codebook, m, subDim) {
  const codes = new Uint8Array(m);
  
  for (let i = 0; i < m; i++) {
    const subVec = vector.slice(i * subDim, (i + 1) * subDim);
    let bestIdx = 0, minDist = Infinity;
    
    // Tìm centroid gần nhất trong codebook của subspace
    for (let cIdx = 0; cIdx < codebook[i].length; cIdx++) {
      const centroid = codebook[i][cIdx];
      const dist = subVec.reduce((sum, val, d) => sum + Math.pow(val - centroid[d], 2), 0);
      if (dist < minDist) {
        minDist = dist;
        bestIdx = cIdx;
      }
    }
    codes[i] = bestIdx; // Ghi chỉ số byte
  }
  return codes;
}

7.4 Tìm kiếm bất đối xứng (Asymmetric Distance Computation - ADC)

Nếu chúng ta phải giải nén (de-quantize) các vector trong database quay ngược lại dạng Float32 để so khớp khoảng cách với vector truy vấn, CPU sẽ mất rất nhiều thời gian giải mã và tốc độ tìm kiếm sẽ rất chậm.

Để giải quyết điều này, chúng ta sử dụng phương pháp Tìm kiếm bất đối xứng (Asymmetric Distance Computation - ADC). Trong ADC, vector truy vấn $\mathbf{q}$ vẫn được giữ nguyên ở dạng Float32 thô (chính xác tuyệt đối), chỉ có các vector lưu trong database là ở dạng mã nén PQ $\mathbf{y}$:

Quy trình ADC tối ưu hóa bằng Look-up Table (LUT) diễn ra như sau:

1. Trước khi quét database, tính khoảng cách từ từng sub-vector của query $\mathbf{q}^i$ đến tất cả $k^*$ centroids trong codebook tương ứng. Lưu các khoảng cách này vào một bảng Look-up Table.
2. Khi quét qua từng vector đã nén trong database, chúng ta không cần tính toán số thực nữa. Chỉ cần dùng mảng byte codes làm chỉ số để tra cứu nhanh (Lookup) điểm số khoảng cách tương ứng trong LUT và cộng dồn lại.

// Tạo bảng Lookup Table (LUT) cho truy vấn ADC
function buildLookupTable(queryVector, codebook, m, subDim) {
  const lut = Array.from({ length: m }, () => []);
  
  for (let i = 0; i < m; i++) {
    const querySub = queryVector.slice(i * subDim, (i + 1) * subDim);
    
    // Tính khoảng cách từ query sub-vector đến tất cả centroids trong subspace
    for (let cIdx = 0; cIdx < codebook[i].length; cIdx++) {
      const centroid = codebook[i][cIdx];
      const dist = querySub.reduce((sum, val, d) => sum + Math.pow(val - centroid[d], 2), 0);
      lut[i].push(dist); // Lưu khoảng cách bình phương vào LUT
    }
  }
  return lut;
}

// Tính khoảng cách ADC siêu nhanh bằng cách tra bảng
function computeAdc(codes, lut, m) {
  let distanceSum = 0;
  for (let i = 0; i < m; i++) {
    const code = codes[i];
    distanceSum += lut[i][code]; // Tra cứu phần tử trong O(1) thay vì nhân ma trận
  }
  return Math.sqrt(distanceSum);
}

7.5 Đánh giá ưu nhược điểm của Product Quantization

Product Quantization là cuộc cách mạng giúp giải quyết bài toán chi phí RAM:

Bảng đối chiếu hiệu năng các phương pháp nén dữ liệu:

7.6 Thực hành: Trực quan hóa nén Product Quantization

Demo dưới đây huấn luyện (train) một ProductQuantization thật (engine vdb-engine.js) với $m=2$ subspaces trên 40 vector 4 chiều ngẫu nhiên. Vì mỗi subspace chỉ còn 2 chiều nên có thể vẽ trực tiếp lên mặt phẳng: mỗi ô màu là vùng ảnh hưởng của một centroid trong codebook, các chấm nhạt là vector con huấn luyện được gán (quantize) về centroid gần nhất. Bấm nút bên dưới để sinh một vector truy vấn 4 chiều mới và xem PQ nén nó thành 2 bytes như thế nào:

import { ProductQuantization } from './vdb-engine.js';

// m = 2 subspaces, moi subspace 2 chieu (vector goc 4 chieu),
// k = 6 centroids moi subspace (nho, de truc quan hoa)
const pq = new ProductQuantization(2, 6, 'euclidean');

// Sinh 40 vector 4 chieu ngau nhien lam du lieu huan luyen
const trainingVectors = Array.from({ length: 40 }, () => [
  Math.random() * 300,
  Math.random() * 300,
  Math.random() * 300,
  Math.random() * 300,
]);

// Train K-means rieng cho tung subspace -> ra 2 codebook doc lap
pq.train(trainingVectors);

// Chen san mot vai vector vao "CSDL" de demo tim kiem ADC
trainingVectors.slice(0, 8).forEach((v, i) => pq.insert(i, v));

// Sinh vector truy van moi va nen no
const query = [Math.random() * 300, Math.random() * 300, Math.random() * 300, Math.random() * 300];
const code = pq.compress(query); // Uint8Array(2) — vi du [3, 5]

// Tim kiem xap xi (Asymmetric Distance Computation) tren du lieu da nen
const results = pq.search(query, 1);

console.log('Ma nen (2 bytes):', Array.from(code));
console.log('Ket qua ADC gan nhat:', results[0]);