Estimator

This part presents how to use the quantized vectors to estimate Euclidean distances and inner product with minimal efforts.

Notation	Description
$\mathbf{o}_r$, $\mathbf{q}_r$	The raw data vector and the raw query vector.
$\mathbf{c}$	The center vector.
$D$	The dimension of vectors.
$B$	The number of bits per dimension used for the quantization.
$c_B$	A constant value $c_B=- \frac{2^B-1}{2}$.
$\mathbf{1}_D$	The all-one vector of dimension $D$.
$P$	The random rotation matrix.
$\mathbf{\bar o}$	The quantized normalized data vector.
$\mathbf{q}_r'$	The reversely rotated raw query vector: $\mathbf{q}_r'=P^{-1}\mathbf{q}_r$.
$S_q$	$S_q=\sum_{i=1}^{D} \mathbf{q}'_{r}[i].$
$\Delta_x$	The rescaling factor of a quantized vector.
$\mathbf{x}_u$	The uint representation of the quantized vectors.

Quantization

Recall that the quantization algorithm of RaBitQ receives a raw data vector $\mathbf{o}_r$, a center vector $\mathbf{c}$ and a sample of random rotation matrices $P$ as inputs and returns a value $\Delta_x$ and an uint vector $\mathbf{x}_u$ as output, such that

$$\begin{align} \mathbf{o}_r - \mathbf{c} \approx \| \mathbf{o}_r-\mathbf{c} \| \cdot \mathbf{\bar o} = \Delta_x P(\mathbf{x}_u+c_B\mathbf{1}_D) \end{align}$$

where $c_B=- \frac{2^B-1}{2}$ and $B$ is the number of bits used for the quantization. $\mathbf{1}_D$ is the all-one vector of dimension $D$. $P$ is a sample of random rotation matrices.

Estimator

The following derivation covers the estimators for Euclidean distances, inner products, and cosine similarity. The cosine similarity is supported by the same estimator as the inner product.

Estimator of Euclidean Distance

$$\begin{align} &\| \mathbf{o}_r-\mathbf{q}_r\|^2 \\=& \|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 - 2 \left< \mathbf{o}_r - \mathbf{c}, \mathbf{q}_r - \mathbf{c}\right> \\ =& \|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 - 2\| \mathbf{o}_r-\mathbf{c}\|\cdot \| \mathbf{q}_r-\mathbf{c}\| \cdot \left<\mathbf{ o},\mathbf{q} \right> \\ \approx& \|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 - 2\| \mathbf{o}_r-\mathbf{c}\|\cdot \| \mathbf{q}_r-\mathbf{c}\| \cdot \frac{\left<\mathbf{\bar o},\mathbf{q} \right>}{ \left<\mathbf{\bar o},\mathbf{o} \right>} \\ =& \|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 - 2\| \mathbf{o}_r-\mathbf{c}\| \cdot \frac{\left<\mathbf{\bar o},\mathbf{q}_r-\mathbf{c} \right>}{\left<\mathbf{\bar o},\mathbf{o} \right>} \\ =& \|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 +2\| \mathbf{o}_r-\mathbf{c}\| \frac{\left< \mathbf{\bar o}, \mathbf{c}\right>}{\left<\mathbf{\bar o},\mathbf{o} \right>} - \frac{2}{\left<\mathbf{\bar o},\mathbf{o} \right>}\cdot \left< \| \mathbf{o}_r-\mathbf{c}\|\cdot P^{-1}\mathbf{\bar o}, \mathbf{q}_r' \right> \\=&\|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 +2\| \mathbf{o}_r-\mathbf{c}\| \frac{\left< \mathbf{\bar o}, \mathbf{c}\right>}{\left<\mathbf{\bar o},\mathbf{o} \right>} - \frac{2\Delta_x}{\left<\mathbf{\bar o},\mathbf{o} \right>}\cdot \left[ \left< \mathbf{x}_u, \mathbf{q}_r' \right>+c_B S_q\right] \\ \| \mathbf{o}_r-\mathbf{q}_r\|^2 \approx&\|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 +2\| \mathbf{o}_r-\mathbf{c}\| \frac{\left< \mathbf{\bar o}, \mathbf{c}\right>}{\left<\mathbf{\bar o},\mathbf{o} \right>} - \frac{2\Delta_x}{\left<\mathbf{\bar o},\mathbf{o} \right>}\cdot \left[ \left< \mathbf{x}_u, \mathbf{q}_r' \right>+c_B S_q\right] \end{align}$$

The error bound is given by

$$\begin{align} % \pm \ 2 \|\mathbf{o}_r-\mathbf{c}\|\cdot \| \mathbf{q}_r-\mathbf{c}\| \cdot \frac{c_{error}}{\sqrt{D}\cdot 2^{B-1}}\\ \pm \ 2 \|\mathbf{o}_r-\mathbf{c}\|\cdot \| \mathbf{q}_r-\mathbf{c}\| \cdot \sqrt{\frac{1 - \left< \mathbf{\bar o},\mathbf{o}\right>^2}{\left< \mathbf{\bar o},\mathbf{o}\right>^2}} \frac{\epsilon_{0}}{\sqrt{D-1}} \end{align}$$

Here $\epsilon_0$ is a tunable parameter which controls the confidence level of the error bound. By default, $\epsilon_0=1.9$ guarantees nearly perfect confidence.

We store the following variables such that the estimator can be computed easily.

Name (Type) of Variable	Description
F_add (float)	$\|\mathbf{o}_r-\mathbf{c}\|^2+2\| \mathbf{o}_r-\mathbf{c}\| \frac{\left< \mathbf{\bar o}, \mathbf{c}\right>}{\left<\mathbf{\bar o},\mathbf{o} \right>}$
F_rescale (float)	$-2\frac{\Delta_x}{\left<\mathbf{\bar o},\mathbf{o} \right>}$
F_error (float)	$\ 2 \|\mathbf{o}_r-\mathbf{c}\| \cdot \sqrt{\frac{1 - \left< \mathbf{\bar o},\mathbf{o}\right>^2}{\left< \mathbf{\bar o},\mathbf{o}\right>^2}} \frac{\epsilon_{0}}{\sqrt{D-1}}$
G_add (float)	$\|\mathbf{q}_r-\mathbf{c}\|^2$
G_kBxSumq (float)	$c_BS_q$
G_error (float)	$\| \mathbf{q}_r-\mathbf{c}\|$

Estimator of Inner Product

When inner product is used as the metric of vector search, it targets the data vector which has the maximum inner product with the query vector. To unify the question with nearest neighbor search, we follows Faiss and hnswlib to compute the negative inner product.

$$\begin{align} &-\left< \mathbf{o}_r,\mathbf{q}_r\right> \\=& -\left< \mathbf{o}_r-\mathbf{c} + \mathbf{c},\mathbf{q}_r-\mathbf{c} + \mathbf{c}\right> \\=& -\left< \mathbf{q}_r,\mathbf{c}\right> -\left< \mathbf{o}_r-\mathbf{c},\mathbf{c}\right> - \left< \mathbf{o}_r-\mathbf{c},\mathbf{q}_r-\mathbf{c} \right> \\ \approx &-\left< \mathbf{q}_r,\mathbf{c}\right> -\left< \mathbf{o}_r-\mathbf{c},\mathbf{c}\right> +\| \mathbf{o}_r-\mathbf{c}\| \frac{\left< \mathbf{\bar o}, \mathbf{c}\right>}{\left<\mathbf{\bar o},\mathbf{o} \right>} - \frac{\Delta_x}{\left<\mathbf{\bar o},\mathbf{o} \right>}\cdot \left[ \left< \mathbf{x}_u, \mathbf{q}_r' \right>+c_B S_q\right] \\ -\left< \mathbf{o}_r,\mathbf{q}_r\right>\approx &-\left< \mathbf{q}_r,\mathbf{c}\right> -\left< \mathbf{o}_r-\mathbf{c},\mathbf{c}\right> +\| \mathbf{o}_r-\mathbf{c}\| \frac{\left< \mathbf{\bar o}, \mathbf{c}\right>}{\left<\mathbf{\bar o},\mathbf{o} \right>} - \frac{\Delta_x}{\left<\mathbf{\bar o},\mathbf{o} \right>}\cdot \left[ \left< \mathbf{x}_u, \mathbf{q}_r' \right>+c_B S_q\right] \end{align}$$

The error bound is given by

$$\begin{align} % \pm \|\mathbf{o}_r-\mathbf{c}\|\cdot \| \mathbf{q}_r-\mathbf{c}\| \cdot \frac{c_{error}}{\sqrt{D}\cdot 2^{B-1}} \pm \|\mathbf{o}_r-\mathbf{c}\|\cdot \| \mathbf{q}_r-\mathbf{c}\| \cdot \sqrt{\frac{1 - \left< \mathbf{\bar o},\mathbf{o}\right>^2}{\left< \mathbf{\bar o},\mathbf{o}\right>^2}} \frac{\epsilon_{0}}{\sqrt{D-1}} \end{align}$$

Here $\epsilon_0$ is a tunable parameter which controls the confidence level of the error bound. By default, $\epsilon_0=1.9$ guarantees nearly perfect confidence.

We store the following variables such that the estimator can be computed easily. The variables F_ indicates the factors for the data vector. The variables G_ indicates the factors for the query vector.

Name (Type) of Variable	Description
F_add (float)	$-\left< \mathbf{o}_r-\mathbf{c},\mathbf{c}\right>+\| \mathbf{o}_r-\mathbf{c}\| \frac{\left< \mathbf{\bar o}, \mathbf{c}\right>}{\left<\mathbf{\bar o},\mathbf{o} \right>}$
F_rescale (float)	$-\frac{\Delta_x}{\left<\mathbf{\bar o},\mathbf{o} \right>}$
F_error (float)	$\|\mathbf{o}_r-\mathbf{c}\| \cdot \sqrt{\frac{1 - \left< \mathbf{\bar o},\mathbf{o}\right>^2}{\left< \mathbf{\bar o},\mathbf{o}\right>^2}} \frac{\epsilon_{0}}{\sqrt{D-1}}$
G_add (float)	$-\left< \mathbf{q}_r,\mathbf{c}\right>$
G_kBxSumq (float)	$c_BS_q$
G_k1xSumq (float)	$c_1S_q$
G_error (float)	$\| \mathbf{q}_r-\mathbf{c}\|$

Implementation

Distance Estimation

Based on the above derivation, we can implement the estimator of Euclidean distance and inner product using exactly the same implementation. A pseudo code is given below. Let ip be the inner product between the binary code and the randomly rotated query vector.

// Compute the estimated distance
// Note that G_add is dependent on the center vector.
float est_dist = F_add + G_add + F_rescale * (ip + G_kBxSumq)

// Compute the error bound 
// Note that G_error is dependent on the center vector.
float error_bound = F_error * G_error

// Compute the lower and upper bounds of the estimated distance
float lb_dist = est_dist - error_bound
float ub_dist = est_dist + error_bound

Incremental Distance Estimation

RaBitQ supports incremental distance estimation. We split the code into two parts: the binary code (the most significant bits) and the extended code (the remaining $B-1$ bits). Incremental distance estimation supports to first estimate a coarse distance based on the binary code. If the accuracy is insufficient, we then access the extended code to boost the accuracy. For this, we need to prepare factors for both the binary code and the extended code. The factors for the binary code are stored in F_add, F_rescale and F_error. The factors for the extended code are stored in F_add_ex, F_rescale_ex and F_error_ex. The factors for the query includes G_add, G_error, G_kBxSumq and G_k1xSumq. Let ip_bin be the inner product between the binary code and the randomly rotated query vector and ip_ex be the inner product between the ex-code and the randomly rotated query vector.

// 1-bit dist
float est_dist = F_add + G_add + F_rescale * (ip_bin + G_k1xSumq)
float bound = F_error * G_error
float ub_dist = est_dist + bound
float lb_dist = est_dist - bound

// boost to full-bit dist
float ex_est_dist = F_add_ex + G_add + F_rescale_ex * (ip_bin << (bits - 1) + ip_ex + G_kBxSumq)
float ex_bound = F_error_ex * G_error
float ex_ub_dist = ex_est_dist + ex_bound;
float ex_lb_dist = ex_est_dist - ex_bound;