quant

package
v0.1.0 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Jun 20, 2026 License: MIT Imports: 4 Imported by: 0

Documentation

Overview

Package quant implements TurboQuant, a data-oblivious vector quantizer with near-optimal distortion (Zandieh, Daliri, Hadian, Mirrokni, 2025; arXiv:2504.19874).

The construction has three parts:

  • A randomized Hadamard rotation spreads a vector's energy evenly across coordinates, so the coordinates of a unit vector become approximately i.i.d. Gaussian. The rotation is orthonormal and preserves inner products exactly, which is what lets a query stay in full precision while the database is quantized (asymmetric distance computation).
  • An optimal per-coordinate scalar quantizer, the Lloyd-Max codebook for the standard normal, encodes the rotated coordinates. The norm is stored separately.
  • A 1-bit QJL sketch of the quantization residual corrects the bias that the MSE-optimal quantizer would otherwise introduce into inner-product estimates.

Two families of estimators are exposed because they trade off differently. The Score family is the low-variance main term, best for ranking and candidate generation. The IP/L2/Cosine family adds the residual correction for unbiased magnitudes at the cost of the sketch's variance. See the Query type for details.

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

This section is empty.

Types

type BenchOptions

type BenchOptions struct {
	N        int    // database vectors (default 5000)
	Queries  int    // query vectors (default 100)
	TopK     int    // recall cutoff (default 10)
	Clusters int    // cluster centers, for realistic structure (default 16)
	Seed     uint64 // determinism
}

BenchOptions configures a synthetic benchmark of a quantizer configuration. Zero fields take sensible defaults.

type BenchResult

type BenchResult struct {
	Dim, Bits, ResidualDims int
	CodeBytes               int     // bytes to store one encoded vector (packed)
	CompressionRatio        float64 // float32 input bytes / code bytes
	RecallAtK               float64 // quantized Score ranking vs exact float ranking
	CodebookMSE             float64 // textbook Lloyd-Max distortion of the codebook
	EncodeVecsPerSec        float64
	QueryScoresPerSec       float64
}

BenchResult reports the accuracy, size, and speed tradeoff of a quantizer configuration. The accuracy and size fields are deterministic for a fixed config and seed; the throughput fields depend on the host.

func Benchmark

func Benchmark(cfg Config, opt BenchOptions) BenchResult

Benchmark builds a quantizer for cfg and measures, on synthetic clustered unit vectors: recall@TopK of the low-variance Score ranking against the exact float ranking, the packed code size and compression ratio, and encode/query throughput. It is the basis of the `turbograph quant bench` tool and is usable directly as a library call.

type Code

type Code struct {
	Codes   []uint8 // per-coordinate codebook indices, length Dim()
	Norm    float32 // original Euclidean norm
	ResNorm float32 // residual norm in the standardized (z) space
	Signs   uint64  // QJL residual sign bits, one per projection
}

Code is the compressed representation of one vector.

type Codebook

type Codebook struct {
	// contains filtered or unexported fields
}

Codebook is an optimal scalar quantizer for the standard normal distribution, computed by the Lloyd-Max iteration. After a random rotation the coordinates of a unit vector are close to i.i.d. N(0, 1/D); scaled by sqrt(D) they are close to N(0,1), which is exactly the distribution this codebook is tuned for.

For the optimal (centroid) quantizer the reconstruction level of a cell is the conditional mean of the source inside that cell, which yields two identities used throughout the estimators:

E[Q(X)^2] = E[X*Q(X)]            (reconstruction is the MMSE estimate)
MSE = E[(X-Q(X))^2] = 1 - E[Q(X)^2]

func (*Codebook) Bits

func (c *Codebook) Bits() int

Bits returns the number of bits per coordinate.

func (*Codebook) Decode

func (c *Codebook) Decode(code uint8) float32

Decode returns the reconstruction value for a code.

func (*Codebook) Encode

func (c *Codebook) Encode(z float32) uint8

Encode maps a standard-normal-scaled value to its codebook index.

func (*Codebook) Levels

func (c *Codebook) Levels() []float32

Levels returns the reconstruction values (read-only).

func (*Codebook) MSE

func (c *Codebook) MSE() float64

MSE returns the expected squared quantization error against N(0,1).

type Config

type Config struct {
	Dim          int    // input vector dimension
	Bits         int    // bits per coordinate, 1..8
	Rounds       int    // rotation mixing passes (default 3)
	ResidualDims int    // QJL residual projections m, 0..64; 0 disables debiasing
	Seed         uint64 // determinism seed
}

Config parameterizes a Quantizer.

type PCG

type PCG struct {
	// contains filtered or unexported fields
}

PCG is a small, fast, deterministic pseudo-random generator (xoshiro256** seeded through splitmix64). It exists so the data-oblivious randomness used by the quantizer (sign flips, JL projections) is reproducible from a single seed and independent of the global math/rand state.

func NewPCG

func NewPCG(seed uint64) *PCG

NewPCG seeds a generator from a 64-bit seed.

func (*PCG) Float64

func (p *PCG) Float64() float64

Float64 returns a value uniformly distributed in [0, 1).

func (*PCG) NormFloat64

func (p *PCG) NormFloat64() float64

NormFloat64 returns a standard normal sample via the Box-Muller transform.

func (*PCG) Uint64

func (p *PCG) Uint64() uint64

Uint64 returns the next 64-bit value.

type Quantizer

type Quantizer struct {
	// contains filtered or unexported fields
}

Quantizer implements TurboQuant: a data-oblivious vector quantizer that rotates inputs to an isotropic frame, applies an optimal per-coordinate scalar quantizer, and (optionally) captures a 1-bit QJL sketch of the quantization residual so inner products can be estimated without bias.

A single Quantizer is immutable after construction and safe for concurrent use by encoders and queries.

func New

func New(cfg Config) *Quantizer

New builds a Quantizer from cfg.

func (*Quantizer) Bits

func (q *Quantizer) Bits() int

Bits returns bits per coordinate.

func (*Quantizer) Codebook

func (q *Quantizer) Codebook() *Codebook

Codebook exposes the underlying scalar quantizer (read-only use).

func (*Quantizer) Decode

func (q *Quantizer) Decode(c Code) []float32

Decode reconstructs an approximation of the original vector (in the original, unrotated frame, truncated to the input dimension).

func (*Quantizer) Dim

func (q *Quantizer) Dim() int

Dim returns the padded working dimension.

func (*Quantizer) Encode

func (q *Quantizer) Encode(x []float32) Code

Encode compresses a vector. The returned Code owns its Codes slice.

func (*Quantizer) EncodeBatch

func (q *Quantizer) EncodeBatch(rows []float32, n, stride int) []Code

EncodeBatch encodes many vectors with shared scratch buffers. rows is a flat row-major matrix with stride equal to the input dimension.

func (*Quantizer) EncodeInto

func (q *Quantizer) EncodeInto(x, buf, res []float32, c *Code)

EncodeInto encodes x reusing caller-provided scratch buffers (buf and res must have length Dim(), and c.Codes must too). It performs no allocation, which matters when encoding large corpora.

func (*Quantizer) PrepareQuery

func (q *Quantizer) PrepareQuery(vec []float32) *Query

PrepareQuery rotates q and builds its lookup tables.

func (*Quantizer) Rotation

func (q *Quantizer) Rotation() *Rotation

Rotation exposes the underlying rotation (read-only use).

type Query

type Query struct {
	// contains filtered or unexported fields
}

Query is a prepared query vector. Building it once amortizes the rotation and the per-coordinate lookup-table construction across a whole index scan.

The main inner-product term is evaluated as an asymmetric distance computation (ADC): for code c, sum_i qlut[i*k + c_i] is exactly sum_i rotated_query_i * level(c_i). The query stays in full precision; only the database is quantized.

Two estimators are exposed because they have different bias/variance tradeoffs:

  • Score is the low-variance main term. It is slightly biased in scale but order-preserving, which makes it the right choice for ranking and candidate generation in nearest-neighbor search.
  • IP/L2/Cosine add the 1-bit QJL residual correction (when the quantizer was built with ResidualDims > 0). That removes the bias so absolute magnitudes are accurate, at the cost of the variance the sketch introduces.

func (*Query) Cosine

func (qr *Query) Cosine(c Code) float32

Cosine estimates the cosine similarity between the query and x.

func (*Query) CosineScore

func (qr *Query) CosineScore(c Code) float32

CosineScore returns a low-variance, order-preserving score whose order matches cosine similarity. It omits the query-norm factor (constant across a scan), so it is proportional to cosine and ideal for ranking by direction.

func (*Query) IP

func (qr *Query) IP(c Code) float32

IP returns an unbiased estimate of the inner product <query, x>.

func (*Query) L2

func (qr *Query) L2(c Code) float32

L2 returns an unbiased estimate of the squared Euclidean distance ||query - x||^2.

func (*Query) L2Score

func (qr *Query) L2Score(c Code) float32

L2Score returns a low-variance, order-preserving score whose ascending order matches squared Euclidean distance. Use it for nearest-neighbor ranking.

func (*Query) Score

func (qr *Query) Score(c Code) float32

Score returns a low-variance, order-preserving estimate of <query, x>. Use it for ranking and candidate generation; it is the fastest path and the most accurate for ordering.

type Rotation

type Rotation struct {
	// contains filtered or unexported fields
}

Rotation is an orthonormal, data-oblivious transform built from alternating Rademacher sign flips and fast Walsh-Hadamard transforms (an SRHT-style operator). Applied to a vector it spreads energy evenly across coordinates so that, for a unit input, each output coordinate behaves like an independent zero-mean Gaussian with variance 1/D. That is the property TurboQuant relies on to make per-coordinate scalar quantization near optimal.

The transform is orthonormal, so it preserves both norms and inner products exactly. The same Rotation must be applied to database and query vectors for asymmetric distance estimation to be consistent.

func NewRotation

func NewRotation(inDim, rounds int, rng *PCG) *Rotation

NewRotation builds a rotation for inputs of dimension inDim. Vectors are zero-padded to the next power of two. rounds controls mixing quality; three passes is sufficient to make coordinates close to Gaussian for any input.

func (*Rotation) Apply

func (r *Rotation) Apply(dst, src []float32)

Apply writes the rotated image of src into dst. src may be shorter than the padded dimension; it is zero-extended. dst must have length Dim().

func (*Rotation) Dim

func (r *Rotation) Dim() int

Dim returns the padded working dimension (a power of two).

func (*Rotation) InputDim

func (r *Rotation) InputDim() int

InputDim returns the original, unpadded input dimension.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL