TurboQuant: Add distortion benchmark

[connortsui20]

Summary

Tracking issue: #7830

Adds distortion / error benchmarking option for TurboQuant. Note that this benchmark is still over the old implementation of TurboQuant, but because the algorithm is identical it shouldn't affect any numbers.

We will cut over to the new version atomically once everything is implemented in vortex-turboquant.

Testing

[codspeed-hq]

Merging this PR will not alter performance

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

✅ 1266 untouched benchmarks

---

_{Comparing ct/tq-error (11b5de1) with develop (a2323f1)}

[connortsui20]

whoops forgot to make sure the binary is always up to date (edit: fixed)

[danking]

@claude review with careful attention paid to the validity of the statistical arguments.

[danking]

I hid ChatGPT's comment because it is quite long and I think not quite correct about some statistical things, but you may find it helpful to read over it if anything I've written below is confusing.

[danking]

I think we need to switch to the correct definition of mean square error, $\frac{1}{N} \sum_i (x_i - \hat{x_i})^2$. This also means we don't need to elide low magnitude rows, because MSE is well defined regardless of norm.

[connortsui20]

ok so I decided to just remove the "normalization" factor of the dimensions since higher dimensions should theoretically result in lower distortion anyways, so not its just sum of squared errors.

[connortsui20]

oh wait I need to make these all unit-norm first

[danking]

Claude mentions this as well but I think this section has a few errors.

Equation 2 from the TurboQuant paper defines the inner product error as: for all pair of vectors, x and y, in $R^d$, the following quantity

is upper bounded by (from an unnumbered equation under "Inner Product TurboQuant"):

We need to respect the quantifier and expectation ordering. We need to fix a bit width, an x vector, and a y vector, and then compute the squared cosine distance:

d_prods = []
for seed in random_seeds:
    diff = compute_cosines(y, x) - compute_cosines(y, decode(encode(x, seed)))
    d_prods.append(diff * diff)

d_prod_mean = sum(d_prods) / len(d_prods)

That value, d_prod_mean, in the limit of looping over all seeds, is bound by the error.

[danking]

I'm not sure how to choose y. I think choosing a set of random vectors on startup is reasonable. The property should hold for any y you choose.

Again, another unnumbered equation, this one under "1.1 Problem Definition", but this argues that if we choose y = x, then the expectation of the difference is just zero. That seems like a separately useful statistic to measure and display.

[connortsui20]

I think it is fine to just have a bunch of random pairs of x and y, I think a specific random y and then a random set of x is fine in the expectation calculation because everything (should be) independent.

[connortsui20]

issue 1 is not a real issue because I say "normalized" in the chart but I adjusted it anyways, issue 2 is fake, and I just fixed issue 3 and 4

[danking]

There are still some issues with the statistics. I'm happy to jump on a zoom tomorrow morning to talk through it.

[danking]

This equality is only true if the norm of x is equal to the norm of its reconstruction. I think we should just elide the equality.

[danking]

A "per-vector NMSE" doesn't make sense to me. NSME is "normalized mean square error", right? I agree that $||x - x'||^2 / ||x||^2$ is a normalized squared error, but there's no mean/expectation. I think the correct verbiage is:

Suggested change

	//! Reports per-vector NMSE (`||x - x'||^2 / ||x||^2 = ||unit(x) - unit(x')||^2`) and per-
	//! Reports the normalized mean square error of this TurboQuant implementation over vectors from a configurable dataset [...]

No "per-vector" stuff.

[danking]

This gives me heebie jeebies. Any zero vector in the test data erroneously reduces our estimate of the reconstruction error.

Suppose, for example, if all my test vectors are the zero vector and all my reconstructions are the all ones vector. The reconstruction error, as defined by the paper, would be $\sqrt{d}$! That's huge! But this function reports the reconstruction error as zero!

Are there any zero vectors in the test data?

[danking]

Actually, why not just normalize the vectors after loading them? Then everything downstream can follow the paper's formalism exactly (e.g. we can just use mean square error here instead of normalizing). If any vectors are zero vectors, we just throw them out while loading.

[danking]

This is still not equivalent to $D_{mse}$ from the paper. This is, given some specific seed,

$\mathbb{E}_{x \in \mathbb{R}^d}[||x - Q^{-1}(Q(x, \textrm{seed}))||^2]$

The paper defines $D_{mse}$ as, for some specific vector x:

$\mathbb{E}_{\mathrm{seed} \in R}[||x - Q^{-1}(Q(x, \textrm{seed}))||^2]$

I realize you're trying to replicate Figure 3. I think we can do that, but it's a less useful figure than Figure 1. Figure 1 shows the distribution of $D_{prod}$ over multiple bit widths and over all pairings of the 100,000 data vectors and the 1,000 query vectors. I'm not sure why the paper lacks a similar figure for $D_{mse}$. That seems to me a very reasonable figure to generate. Regardless, if you want to construct Figure 3, it should be constructed according to this:

$E_{x \in \mathbb{R}^d} [ \mathbb{E}_{\mathrm{seed} \in R}[||x - Q^{-1}(Q(x, \textrm{seed}))||^2]]$

Expectations are linear operators so your for loops can be for seed: for x: or for x: for seed:; however, there you must average over multiple seeds for your plots to fairly characterize Vortex's TurboQuant.