234 lines
8.6 KiB
C++
234 lines
8.6 KiB
C++
// Catch2 unit tests for the TurboQuantProd mathematical contract
|
||
// Tests the QJL gamma / sign-agreement / arcsin relationship that underpins
|
||
// estimateInnerProductFull(). Validates the Joint Normal Lemma result:
|
||
//
|
||
// E[sgn(S·r1)_m · sgn(S·r2)_m] = (2/π) · arcsin(ρ_res)
|
||
//
|
||
// where ρ_res = (r1·r2) / (||r1||·||r2||) is the residual correlation.
|
||
//
|
||
// NOTE: This file tests the QJL mathematical foundation independently of the
|
||
// TurboQuantProd implementation. It verifies that our interpretation of
|
||
// gamma (raw agreement rate) in the estimator is mathematically sound.
|
||
|
||
#include <catch2/catch_test_macros.hpp>
|
||
|
||
#include <algorithm>
|
||
#include <cmath>
|
||
#include <random>
|
||
#include <vector>
|
||
#include <yams/vector/turboquant.h>
|
||
|
||
using namespace yams::vector;
|
||
|
||
namespace {
|
||
|
||
constexpr double kPi = 3.14159265358979323846;
|
||
|
||
// Verify E[sgn(a)·sgn(b)] = (2/π)·arcsin(ρ) for Gaussian pairs
|
||
// by generating many trials and comparing empirical mean to theoretical value.
|
||
TEST_CASE("QJL gamma contract: E[sgn·sgn] = (2/π)·arcsin(ρ)", "[turboquant][math][qjl][catch2]") {
|
||
std::mt19937 rng(137);
|
||
|
||
// Test several correlation values: 0 (orthogonal), 0.3, 0.7, 1.0 (identical)
|
||
std::vector<double> rhos = {0.0, 0.3, 0.7, 1.0};
|
||
constexpr size_t kTrials = 10000;
|
||
|
||
for (double rho : rhos) {
|
||
double sum_sign_product = 0.0;
|
||
for (size_t t = 0; t < kTrials; ++t) {
|
||
// Generate two correlated standard normal variables (X, Y) with correlation rho
|
||
// X ~ N(0,1), Y = rho·X + sqrt(1-rho²)·Z, Z~N(0,1) independent of X
|
||
std::normal_distribution<float> dist(0.0f, 1.0f);
|
||
float x = dist(rng);
|
||
float z = dist(rng);
|
||
float rhoFloat = static_cast<float>(rho);
|
||
float y = rhoFloat * x + std::sqrt(1.0f - rhoFloat * rhoFloat) * z;
|
||
|
||
float sx = (x >= 0.0f) ? 1.0f : -1.0f;
|
||
float sy = (y >= 0.0f) ? 1.0f : -1.0f;
|
||
float sp = sx * sy; // sgn(a)·sgn(b)
|
||
|
||
sum_sign_product += sp;
|
||
}
|
||
|
||
double empirical = sum_sign_product / static_cast<double>(kTrials);
|
||
double theoretical = (2.0 / kPi) * std::asin(rho);
|
||
|
||
INFO("rho=" << rho << ": empirical=" << empirical << ", theoretical=" << theoretical
|
||
<< ", diff=" << (empirical - theoretical));
|
||
|
||
// Allow generous tolerance due to variance in finite samples (10% or 0.05 absolute)
|
||
double abs_diff = std::abs(empirical - theoretical);
|
||
CHECK(abs_diff < std::max(0.05, std::abs(theoretical) * 0.15));
|
||
}
|
||
}
|
||
|
||
// Verify the inverse: arcsin(E[sgn·sgn] · π/2) ≈ ρ
|
||
// This is what estimateInnerProductFull() does internally.
|
||
TEST_CASE("QJL gamma contract: arcsin inversion recovers rho", "[turboquant][math][qjl][catch2]") {
|
||
std::mt19937 rng(911);
|
||
constexpr size_t kTrials = 10000;
|
||
|
||
// Test with a fixed rho but many trials
|
||
double rho = 0.5;
|
||
double sum_sign_product = 0.0;
|
||
|
||
for (size_t t = 0; t < kTrials; ++t) {
|
||
std::normal_distribution<float> dist(0.0f, 1.0f);
|
||
float x = dist(rng);
|
||
float z = dist(rng);
|
||
float rhoFloat = static_cast<float>(rho);
|
||
float y = rhoFloat * x + std::sqrt(1.0f - rhoFloat * rhoFloat) * z;
|
||
|
||
float sx = (x >= 0.0f) ? 1.0f : -1.0f;
|
||
float sy = (y >= 0.0f) ? 1.0f : -1.0f;
|
||
sum_sign_product += sx * sy;
|
||
}
|
||
|
||
double empirical_sp = sum_sign_product / static_cast<double>(kTrials);
|
||
double recovered_rho = std::sin((kPi / 2.0) * empirical_sp);
|
||
|
||
INFO("Input rho=" << rho << ", empirical E[sgn·sgn]=" << empirical_sp
|
||
<< ", recovered rho=" << recovered_rho);
|
||
CHECK(std::abs(recovered_rho - rho) < 0.05);
|
||
}
|
||
|
||
// Verify the sign agreement rate gamma_raw vs sign product relationship:
|
||
// gamma_raw = (1 + E[sgn·sgn]) / 2
|
||
// For orthogonal vectors (E[sgn·sgn]=0): gamma_raw ≈ 0.5
|
||
TEST_CASE("QJL gamma contract: orthogonal vectors give gamma_raw ≈ 0.5",
|
||
"[turboquant][math][qjl][catch2]") {
|
||
std::mt19937 rng(777);
|
||
constexpr size_t kTrials = 50000;
|
||
|
||
double sum_sign_product = 0.0;
|
||
for (size_t t = 0; t < kTrials; ++t) {
|
||
std::normal_distribution<float> dist(0.0f, 1.0f);
|
||
float x = dist(rng);
|
||
float y = dist(rng); // Independent of x
|
||
|
||
float sx = (x >= 0.0f) ? 1.0f : -1.0f;
|
||
float sy = (y >= 0.0f) ? 1.0f : -1.0f;
|
||
sum_sign_product += sx * sy;
|
||
}
|
||
|
||
double empirical_sp = sum_sign_product / static_cast<double>(kTrials);
|
||
double gamma_raw = (1.0 + empirical_sp) / 2.0;
|
||
|
||
INFO("Orthogonal Gaussian pair: E[sgn·sgn]=" << empirical_sp << ", gamma_raw=" << gamma_raw);
|
||
CHECK(std::abs(empirical_sp) < 0.03); // Should be close to 0
|
||
CHECK(std::abs(gamma_raw - 0.5) < 0.015); // gamma_raw should be close to 0.5
|
||
}
|
||
|
||
// Correct inversion: arcsin(rho) = pi*gamma - pi/2 -> rho = sin(pi*gamma - pi/2) = -cos(pi*gamma)
|
||
// For orthogonal (rho=0): gamma_raw=0.5, rho_res=sin(pi/2-pi/2)=sin(0)=0
|
||
// For identical (rho=1): gamma_raw=1.0, rho_res=sin(pi-pi/2)=sin(pi/2)=1
|
||
TEST_CASE("QJL gamma contract: raw agreement maps to rho via -cos(pi*gamma)",
|
||
"[turboquant][math][qjl][catch2]") {
|
||
std::mt19937 rng(333);
|
||
constexpr size_t kTrials = 50000;
|
||
|
||
// Compute empirical gamma_raw from many orthogonal Gaussian pairs
|
||
size_t matching = 0;
|
||
for (size_t t = 0; t < kTrials; ++t) {
|
||
std::normal_distribution<float> dist(0.0f, 1.0f);
|
||
float x = dist(rng);
|
||
float y = dist(rng);
|
||
float sx = (x >= 0.0f) ? 1.0f : -1.0f;
|
||
float sy = (y >= 0.0f) ? 1.0f : -1.0f;
|
||
if (sx == sy)
|
||
matching++;
|
||
}
|
||
|
||
double gamma_raw = static_cast<double>(matching) / kTrials;
|
||
// Correct formula: rho_res = -cos(pi*gamma) = sin(pi*gamma - pi/2)
|
||
double rho_res = -std::cos(kPi * gamma_raw);
|
||
|
||
INFO("Orthogonal: gamma_raw=" << gamma_raw << ", rho_res=" << rho_res);
|
||
// For orthogonal: gamma_raw~0.5, rho_res = -cos(pi*0.5) = -cos(pi/2) = 0
|
||
CHECK(gamma_raw > 0.48);
|
||
CHECK(gamma_raw < 0.52); // gamma_raw ~ 0.5
|
||
CHECK(rho_res > -0.1); // rho_res ~ 0 (orthogonal)
|
||
CHECK(rho_res < 0.1);
|
||
}
|
||
|
||
// Verify that TurboQuantEncoded::encode() captures non-zero residual norms
|
||
// for non-trivially quantized vectors.
|
||
TEST_CASE("TurboQuantEncoded: encode() populates residual_norm_sq for unit vectors",
|
||
"[turboquant][math][turboquant_prod][catch2]") {
|
||
TurboQuantConfig config;
|
||
config.dimension = 128;
|
||
config.bits_per_channel = 4;
|
||
config.seed = 42;
|
||
config.inner_product_mode = true;
|
||
config.qjl_m = 32;
|
||
|
||
TurboQuantProd quantizer(config);
|
||
std::vector<float> v(128);
|
||
std::mt19937 rng(999);
|
||
std::normal_distribution<float> dist(0.0f, 1.0f);
|
||
float norm_sq = 0.0f;
|
||
for (size_t i = 0; i < 128; ++i) {
|
||
v[i] = dist(rng);
|
||
norm_sq += v[i] * v[i];
|
||
}
|
||
float norm = std::sqrt(norm_sq);
|
||
for (size_t i = 0; i < 128; ++i)
|
||
v[i] /= norm;
|
||
|
||
auto enc = quantizer.encode(v);
|
||
|
||
// Residual norm squared should be non-negative
|
||
CHECK(enc.residual_norm_sq >= 0.0f);
|
||
// And should be meaningfully non-zero (4-bit quantization has some error)
|
||
CHECK(enc.residual_norm_sq > 0.001f);
|
||
CHECK(enc.mse_indices.size() == 128);
|
||
CHECK(enc.qjl_signs.size() == 32);
|
||
}
|
||
|
||
// Verify that estimateInnerProduct() is conservative (doesn't wildly inflate estimate)
|
||
TEST_CASE("estimateInnerProduct: blended estimate stays within plausible bounds",
|
||
"[turboquant][math][turboquant_prod][catch2]") {
|
||
TurboQuantConfig config;
|
||
config.dimension = 256;
|
||
config.bits_per_channel = 4;
|
||
config.seed = 55;
|
||
config.inner_product_mode = true;
|
||
config.qjl_m = 64;
|
||
|
||
TurboQuantProd quantizer(config);
|
||
std::mt19937 rng(222);
|
||
|
||
// Test 20 random pairs
|
||
for (size_t trial = 0; trial < 20; ++trial) {
|
||
std::vector<float> v1(256), v2(256);
|
||
for (size_t i = 0; i < 256; ++i)
|
||
v1[i] = std::normal_distribution<float>(0.0f, 1.0f)(rng);
|
||
for (size_t i = 0; i < 256; ++i)
|
||
v2[i] = std::normal_distribution<float>(0.0f, 1.0f)(rng);
|
||
|
||
// Normalize
|
||
float n1_sq = 0.0f, n2_sq = 0.0f;
|
||
for (size_t i = 0; i < 256; ++i) {
|
||
n1_sq += v1[i] * v1[i];
|
||
n2_sq += v2[i] * v2[i];
|
||
}
|
||
float n1 = std::sqrt(n1_sq);
|
||
float n2 = std::sqrt(n2_sq);
|
||
for (auto& x : v1)
|
||
x /= n1;
|
||
for (auto& x : v2)
|
||
x /= n2;
|
||
|
||
auto enc1 = quantizer.encode(v1);
|
||
auto enc2 = quantizer.encode(v2);
|
||
float est = quantizer.estimateInnerProduct(enc1, enc2);
|
||
|
||
// Inner product of unit vectors should be in [-1, 1]
|
||
// With quantization, allow up to ±3.0 (conservative blend max boost ≈ 1.25x)
|
||
CHECK(est > -3.0f);
|
||
CHECK(est < 3.0f);
|
||
}
|
||
}
|
||
|
||
} // namespace
|