README
¶
vectors/pq
quantizer, err := pq.NewQuantizer(pq.DefaultConfig())
vectors/pq implements product quantization primitives for dense vectors.
It is a vector-space mechanism only. Callers own embedding model identity, training corpus selection, index manifests, recall targets, and invalidation when vectors or codebooks change.
Training rejects ragged and non-finite vectors. Decode and asymmetric-distance
operations reject out-of-range codes; DistanceWithTables retains its
documented panic contract for invalid precomputed inputs. Serialized codebooks
are limited to 256 MiB and 65,536 subspaces, which also bounds allocation count
for hostile headers. Load rejects malformed flags, dimensions, truncated or
non-finite codebooks, and leaves the receiving quantizer unchanged on every
failure.
Documentation
¶
Overview ¶
Package pq implements Product Quantization for efficient vector compression.
Overview ¶
Product Quantization (PQ) is a technique for compressing high-dimensional vectors while preserving approximate distance computations. It works by:
- Splitting the vector into M subspaces (e.g., 1536-dim -> 96 subspaces of 16-dim)
- Training K centroids per subspace using k-means clustering
- Encoding each subspace as a single-byte index (if K=256)
This achieves significant compression:
- Original: 1536 dimensions * 4 bytes = 6,144 bytes per vector
- Compressed: 96 subspaces * 1 byte = 96 bytes per vector
- Compression ratio: 64x
Accuracy ¶
PQ provides approximately 95% accuracy for similarity search at recall@10. The accuracy depends on:
- Number of subspaces (M): More subspaces = better accuracy
- Number of centroids (K): More centroids = better accuracy
- Data distribution: Clustered data compresses better
Usage ¶
Basic usage:
// Configure for 1536-dimensional vectors
config := pq.Config{
NumSubspaces: 96,
NumCentroids: 256,
Dimension: 1536,
}
// Create quantizer
q, err := pq.NewQuantizer(config)
if err != nil {
return err
}
// Train on sample vectors (need at least NumCentroids vectors)
err = q.Train(trainingVectors)
if err != nil {
return err
}
// Encode a vector
codes, err := q.Encode(vector)
// Decode to approximate vector
approx, err := q.Decode(codes)
Asymmetric Distance Computation ¶
For search, PQ uses Asymmetric Distance Computation (ADC):
- Query vector stays uncompressed
- Database vectors are compressed
- Distance is computed between uncompressed query and compressed DB vectors
This provides better accuracy than symmetric distance (both compressed).
For batch search, precompute lookup tables:
// Build lookup tables for fast batch distance computation
tables, err := q.BuildLookupTables(queryVector)
// O(M) distance computation per vector (instead of O(M * SubspaceDim))
for _, codes := range databaseCodes {
dist := q.DistanceWithTables(tables, codes)
}
Persistence ¶
Save and load trained quantizers:
// Save
var buf bytes.Buffer
q.Save(&buf)
// Load
q2 := &pq.Quantizer{}
q2.Load(&buf)
Loads are strict and atomic: malformed or truncated state leaves q2 unchanged. Serialized codebooks larger than 256 MiB or with more than 65,536 subspaces are rejected before allocation.
Configuration Guidelines ¶
For 1536-dimensional vectors (OpenAI text-embedding-3-small):
config := pq.DefaultConfig() // NumSubspaces=96, NumCentroids=256
For 768-dimensional vectors (smaller models):
config := pq.Config{
NumSubspaces: 48, // 768/48 = 16 dims per subspace
NumCentroids: 256,
Dimension: 768,
}
Training requirements:
- Minimum: NumCentroids vectors (256 for default config)
- Recommended: 1000-10000 vectors for good codebook quality
Package pq implements Product Quantization for vector compression.
Product Quantization (PQ) compresses high-dimensional vectors by: 1. Splitting the vector into M subspaces (e.g., 1536-dim -> 96 subspaces of 16-dim) 2. Training K centroids per subspace using k-means clustering 3. Encoding each subspace as a single byte index (if K=256)
This achieves ~16x compression (1536 float32 = 6KB -> 96 bytes) with ~95% accuracy.
Search uses asymmetric distance computation (ADC): the query vector stays uncompressed while database vectors are compressed. This provides better accuracy than symmetric distance while maintaining fast search.
Index ¶
- func KMeans(vectors [][]float32, k int, maxIter int) ([][]float32, error)
- func KMeansPlusPlus(vectors [][]float32, k int) ([][]float32, error)
- func MiniBatchKMeans(vectors [][]float32, k int, batchSize int, maxIter int) ([][]float32, error)
- type Config
- type LookupTables
- type Quantizer
- func (q *Quantizer) AsymmetricDistance(query []float32, codes []byte) (float32, error)
- func (q *Quantizer) BuildLookupTables(query []float32) (LookupTables, error)
- func (q *Quantizer) Config() Config
- func (q *Quantizer) Decode(codes []byte) ([]float32, error)
- func (q *Quantizer) DistanceWithTables(tables LookupTables, codes []byte) float32
- func (q *Quantizer) Encode(vector []float32) ([]byte, error)
- func (q *Quantizer) IsTrained() bool
- func (q *Quantizer) Load(r io.Reader) error
- func (q *Quantizer) Save(w io.Writer) error
- func (q *Quantizer) Train(vectors [][]float32) error
Examples ¶
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
func KMeans ¶
KMeans clusters vectors into k centroids using Lloyd's algorithm. Uses k-means++ initialization for better results.
Parameters:
- vectors: the data points to cluster
- k: number of clusters
- maxIter: maximum iterations (typically 20-50)
Returns the k centroids. Empty clusters are handled by reinitializing from the farthest point.
func KMeansPlusPlus ¶
KMeansPlusPlus initializes k centroids using the k-means++ algorithm. This provides better initial centroids than random selection, leading to faster convergence and better clustering quality.
Algorithm: 1. Choose first centroid uniformly at random 2. For each remaining centroid:
- Compute D(x)^2 = distance to nearest existing centroid
- Choose next centroid with probability proportional to D(x)^2
func MiniBatchKMeans ¶
MiniBatchKMeans performs k-means with mini-batches for large datasets. More memory-efficient than full k-means while providing similar quality.
Parameters:
- vectors: the data points to cluster
- k: number of clusters
- batchSize: number of samples per iteration
- maxIter: maximum iterations
Types ¶
type Config ¶
type Config struct {
// NumSubspaces is the number of subspaces (M).
// The vector dimension must be divisible by this.
// Common values: 96 for 1536-dim, 48 for 768-dim.
NumSubspaces int
// NumCentroids is the number of centroids per subspace (K).
// 256 allows single-byte encoding. Must be <= 256.
NumCentroids int
// Dimension is the original vector dimension.
Dimension int
}
Config defines Product Quantization parameters.
func DefaultConfig ¶
func DefaultConfig() Config
DefaultConfig returns a configuration for 1536-dimensional vectors.
Example ¶
package main
import (
"fmt"
"github.com/dotcommander/reliquary/vector/pq"
)
func main() {
cfg := pq.DefaultConfig()
fmt.Println(cfg.NumSubspaces > 0)
}
Output: true
func (Config) CompressionRatio ¶
CompressionRatio returns the compression ratio achieved by PQ. For 1536-dim vectors with 96 subspaces: 6144 bytes -> 96 bytes = 64x.
func (Config) EstimatedAccuracy ¶
EstimatedAccuracy returns the estimated accuracy based on configuration. Higher NumCentroids and NumSubspaces generally improve accuracy.
func (Config) SubspaceDim ¶
SubspaceDim returns the dimensionality of each subspace.
type LookupTables ¶
type LookupTables [][]float32
LookupTables precomputes distances from query subvectors to all centroids. Shape: [NumSubspaces][NumCentroids]
type Quantizer ¶
type Quantizer struct {
// contains filtered or unexported fields
}
Quantizer handles PQ encoding and decoding.
func NewQuantizer ¶
NewQuantizer creates a new quantizer with the given configuration.
func (*Quantizer) AsymmetricDistance ¶
AsymmetricDistance computes the squared L2 distance between a query vector and a compressed vector without fully decoding it. This is more accurate than symmetric distance (both compressed).
func (*Quantizer) BuildLookupTables ¶
func (q *Quantizer) BuildLookupTables(query []float32) (LookupTables, error)
BuildLookupTables precomputes distance tables for fast batch search. After building, use DistanceWithTables for O(M) distance computation.
func (*Quantizer) DistanceWithTables ¶
func (q *Quantizer) DistanceWithTables(tables LookupTables, codes []byte) float32
DistanceWithTables computes distance using precomputed lookup tables. O(M) operations instead of O(M * SubspaceDim).
codes must come from this quantizer's Encode (each value < NumCentroids); an out-of-range code panics.
func (*Quantizer) Encode ¶
Encode compresses a vector to PQ codes. Returns a byte slice of length NumSubspaces.
Source Files