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 ¶
- type BenchOptions
- type BenchResult
- type Code
- type Codebook
- type Config
- type PCG
- type Quantizer
- func (q *Quantizer) Bits() int
- func (q *Quantizer) Codebook() *Codebook
- func (q *Quantizer) Decode(c Code) []float32
- func (q *Quantizer) Dim() int
- func (q *Quantizer) Encode(x []float32) Code
- func (q *Quantizer) EncodeBatch(rows []float32, n, stride int) []Code
- func (q *Quantizer) EncodeInto(x, buf, res []float32, c *Code)
- func (q *Quantizer) PrepareQuery(vec []float32) *Query
- func (q *Quantizer) Rotation() *Rotation
- type Query
- type Rotation
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]
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 (*PCG) NormFloat64 ¶
NormFloat64 returns a standard normal sample via the Box-Muller transform.
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 (*Quantizer) Decode ¶
Decode reconstructs an approximation of the original vector (in the original, unrotated frame, truncated to the input dimension).
func (*Quantizer) EncodeBatch ¶
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 ¶
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 ¶
PrepareQuery rotates q and builds its lookup tables.
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) CosineScore ¶
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) L2 ¶
L2 returns an unbiased estimate of the squared Euclidean distance ||query - x||^2.
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 ¶
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 ¶
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().