trevon/yams
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 2
yams/tests/unit/vector/binary_quantization_catch2_test.cpp
Trevon c26294a938
 feat(topology): opt-in medoid representative and BQ rotation (#277) 
* test(topology): characterize degree representatives and raw BQ sign bits

Pin the current behaviour before adding opt-in representative and BQ
rotation options:

- connected, louvain and kmeans engines on a skewed cluster (dense
  clique plus a weak-edge hub): medoid hash, score, extra routing
  representatives and roles;
- equal-degree ties go to the smaller document hash;
- the unrotated BinaryQuantizedIndex matches a raw-coordinate sign-bit
  reference, with and without a prefix dimension.

Signed-off-by: trvon <me@trevon.dev>

* feat(topology): add an opt-in medoid representative rule

topology.representative_rule = degree (default) | medoid selects how a
cluster's representative document is chosen. Degree keeps the max
within-cluster weighted degree. Medoid picks the member minimising
sum_j sqrt(max(0, 2(1 - cos))) over the other members.

- one shared helper, applyRepresentativeRule(), is called by the
  connected, louvain and kmeans engines after their unchanged degree
  choice; medoid falls back to it without usable geometry;
- clusters above 256 usable members are searched on a deterministic
  hash-ordered stride sample (kMedoidRepresentativeMaxMembers);
- ties within 1e-9 go to the smaller document hash;
- with medoid, farthest-first extra routing representatives are seeded
  from the medoid instead of the centroid;
- the rule is stored in the formerly reserved first header byte of the
  binary snapshot (0 = degree, so old snapshots keep their meaning) and
  extends the construction fingerprint only when it is not degree;
- search debug stats report topology_representative_rule.

Signed-off-by: trvon <me@trevon.dev>

* feat(topology): add an opt-in randomized Hadamard rotation for BQ routing

search.topology.bq_rotation = none (default) | fwht selects the rotation
applied before the centroid BQ shortlist takes sign bits. Raw-coordinate
sign bits make cos(pi * h / d) a poor estimate on anisotropic embeddings;
after x -> (1/sqrt(n)) H D [x; 0] every bit is a random hyperplane.

- RandomizedHadamardRotation in yams/vector: fixed splitmix64 sign
  diagonal plus an in-place FWHT, zero-padded to a power of two; Simeon's
  FWHT projection works on int32 sketches with row subsampling, so a
  small local transform is used instead;
- BinaryQuantizedIndex::build() takes the rotation and seed, rotates
  stored vectors before the prefix is taken, and rotates every float
  query itself, so centroids and queries always share one seed
  (kDefaultBinaryRotationSeed);
- the routing snapshot cache records bqRotation and bqRotationSeed and
  rebuilds the route index when either differs;
- the routing policy fingerprint changes only for a non-default rotation;
- search debug stats report topology_route_bq_rotation.

Signed-off-by: trvon <me@trevon.dev>

* test(bench): map representative rule and BQ rotation arm knobs

Harness-only env knobs written into the isolated benchmark config.toml,
so arms go through ConfigResolver like the other topology settings:

- YAMS_BENCH_TOPOLOGY_REPRESENTATIVE_RULE=degree|medoid
  -> [topology] representative_rule
- YAMS_BENCH_TOPOLOGY_BQ_ROTATION=none|fwht
  -> [search.topology] bq_rotation, next to bq_candidate_limit and
     bq_prefix_dim

The debug-file setup entry records the requested values; per-query
search_stats carry the effective topology_representative_rule and
topology_route_bq_rotation. Portability allowlist lines are shifted for
the inserted code.

Signed-off-by: trvon <me@trevon.dev>

---------

Signed-off-by: trvon <me@trevon.dev>
2026-10-08 14:30:27 -06:00

378 lines
15 KiB
C++
Raw Permalink Blame History

// Copyright (c) 2025 YAMS Contributors
// SPDX-License-Identifier: GPL-3.0-or-later
#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include <yams/vector/binary_quantization.h>
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <numbers>
#include <span>
#include <utility>
#include <vector>
using yams::vector::BinaryQuantizedIndex;
using yams::vector::BinaryQuantizer;
using yams::vector::BinaryVector;
TEST_CASE("BinaryQuantizer: quantize and Hamming distance", "[vector][bq][catch2]") {
SECTION("Empty vector returns empty BinaryVector") {
std::vector<float> empty;
auto bq = BinaryQuantizer::quantize(empty);
CHECK(bq.empty());
CHECK(bq.dimension == 0);
CHECK(bq.words.empty());
}
SECTION("Sign bit quantization correctly packs words") {
// 4 elements: >= 0 is bit 1, < 0 is bit 0
std::vector<float> v = {1.5F, -0.2F, 0.0F, -10.0F};
auto bq = BinaryQuantizer::quantize(v);
REQUIRE(bq.dimension == 4);
REQUIRE(bq.words.size() == 1);
// Bit 0 = 1, Bit 1 = 0, Bit 2 = 1, Bit 3 = 0 -> binary 0101 = 5
CHECK(bq.words[0] == 5ULL);
}
SECTION("Quantization across 64-bit boundaries") {
std::vector<float> v(130, -1.0F);
v[0] = 1.0F; // word 0, bit 0
v[63] = 1.0F; // word 0, bit 63
v[64] = 1.0F; // word 1, bit 0
v[127] = 1.0F; // word 1, bit 63
v[128] = 1.0F; // word 2, bit 0
auto bq = BinaryQuantizer::quantize(v);
REQUIRE(bq.dimension == 130);
REQUIRE(bq.words.size() == 3);
CHECK(bq.words[0] == (1ULL | (1ULL << 63)));
CHECK(bq.words[1] == (1ULL | (1ULL << 63)));
CHECK(bq.words[2] == 1ULL);
}
SECTION("Identical vectors have Hamming distance 0 and cosine similarity 1.0") {
std::vector<float> v = {0.1F, -0.5F, 1.2F, -3.0F, 0.0F, 2.1F};
auto bq1 = BinaryQuantizer::quantize(v);
auto bq2 = BinaryQuantizer::quantize(v);
CHECK(BinaryQuantizer::hammingDistance(bq1, bq2) == 0);
CHECK(BinaryQuantizer::normalizedHammingDistance(bq1, bq2) == Catch::Approx(0.0F));
CHECK(BinaryQuantizer::estimatedCosineSimilarity(bq1, bq2) == Catch::Approx(1.0F));
}
SECTION("Opposite vectors have maximum Hamming distance and cosine similarity -1.0") {
std::vector<float> v1 = {1.0F, 2.0F, 3.0F, 4.0F};
std::vector<float> v2 = {-1.0F, -2.0F, -3.0F, -4.0F};
auto bq1 = BinaryQuantizer::quantize(v1);
auto bq2 = BinaryQuantizer::quantize(v2);
CHECK(BinaryQuantizer::hammingDistance(bq1, bq2) == 4);
CHECK(BinaryQuantizer::normalizedHammingDistance(bq1, bq2) == Catch::Approx(1.0F));
CHECK(BinaryQuantizer::estimatedCosineSimilarity(bq1, bq2) == Catch::Approx(-1.0F));
}
SECTION("Half differing bits yield orthogonal estimated similarity 0.0") {
std::vector<float> v1 = {1.0F, 1.0F, -1.0F, -1.0F};
std::vector<float> v2 = {1.0F, 1.0F, 1.0F, 1.0F};
auto bq1 = BinaryQuantizer::quantize(v1);
auto bq2 = BinaryQuantizer::quantize(v2);
CHECK(BinaryQuantizer::hammingDistance(bq1, bq2) == 2);
CHECK(BinaryQuantizer::normalizedHammingDistance(bq1, bq2) == Catch::Approx(0.5F));
CHECK(std::abs(BinaryQuantizer::estimatedCosineSimilarity(bq1, bq2)) < 1e-5F);
}
}
TEST_CASE("BinaryQuantizedIndex: build and search", "[vector][bq][index][catch2]") {
SECTION("Validation rejects empty or mismatched inputs") {
std::vector<std::size_t> ids = {1, 2};
std::vector<std::vector<float>> vecs = {{1.0F, 0.0F}}; // mismatched size
auto res = BinaryQuantizedIndex::build(ids, vecs);
CHECK_FALSE(res.has_value());
// Duplicate IDs
ids = {1, 1};
vecs = {{1.0F, 0.0F}, {0.0F, 1.0F}};
res = BinaryQuantizedIndex::build(ids, vecs);
CHECK_FALSE(res.has_value());
// Mismatched dimensions
ids = {1, 2};
vecs = {{1.0F, 0.0F}, {1.0F, 0.0F, 1.0F}};
res = BinaryQuantizedIndex::build(ids, vecs);
CHECK_FALSE(res.has_value());
}
SECTION("Search ranks by Hamming distance and respects topK") {
std::vector<std::size_t> ids = {10, 20, 30, 40};
std::vector<std::vector<float>> vecs = {
{1.0F, 1.0F, 1.0F, 1.0F}, // Identical signs: dist 0
{1.0F, 1.0F, 1.0F, -1.0F}, // 1 bit flipped: dist 1
{1.0F, 1.0F, -1.0F, -1.0F}, // 2 bits flipped: dist 2
{-1.0F, -1.0F, -1.0F, -1.0F} // 4 bits flipped: dist 4
};
auto indexRes = BinaryQuantizedIndex::build(ids, vecs);
REQUIRE(indexRes.has_value());
auto index = std::move(indexRes.value());
REQUIRE(index->size() == 4);
REQUIRE(index->dimension() == 4);
CHECK(index->memoryUsageBytes() > 0);
std::vector<float> query = {2.0F, 0.5F, 1.0F, 3.0F}; // all positive signs
auto hits = index->search(query, 3);
REQUIRE(hits.size() == 3);
CHECK(hits[0].id == 10);
CHECK(hits[0].hammingDistance == 0);
CHECK(hits[0].estimatedSimilarity == Catch::Approx(1.0F));
CHECK(hits[1].id == 20);
CHECK(hits[1].hammingDistance == 1);
CHECK(hits[2].id == 30);
CHECK(hits[2].hammingDistance == 2);
}
SECTION("Matryoshka prefix dimension slicing") {
// 128-dimensional vectors, indexed with maxPrefixDimension = 64
std::vector<std::size_t> ids = {1, 2};
std::vector<float> v1(128, 1.0F);
std::vector<float> v2(128, -1.0F);
for (std::size_t i = 64; i < 128; ++i) {
v1[i] = -1.0F;
}
std::vector<std::vector<float>> vecs = {v1, v2};
auto fullIndex = BinaryQuantizedIndex::build(ids, vecs, 0).value();
CHECK(fullIndex->dimension() == 128);
auto prefixIndex = BinaryQuantizedIndex::build(ids, vecs, 64).value();
CHECK(prefixIndex->dimension() == 64);
// Full 128-D query queries 64-D prefix index seamlessly
std::vector<float> query(128, 1.0F);
auto hits = prefixIndex->search(query, 2);
REQUIRE(hits.size() == 2);
CHECK(hits[0].id == 1);
CHECK(hits[0].hammingDistance == 0);
CHECK(hits[1].id == 2);
CHECK(hits[1].hammingDistance == 64);
}
}
namespace {
// Deterministic, platform-independent pseudo-random floats in [-1, 1).
struct SplitMixFloats {
std::uint64_t state;
float next() {
std::uint64_t z = (state += 0x9E3779B97F4A7C15ULL);
z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
z ^= z >> 31;
return static_cast<float>(static_cast<double>(z >> 11) * 0x1.0p-53 * 2.0 - 1.0);
}
};
// Reference raw-coordinate sign-bit shortlist: the behaviour of the unrotated index.
std::vector<std::pair<std::size_t, std::size_t>>
rawSignShortlist(const std::vector<std::vector<float>>& vectors, const std::vector<float>& query,
std::size_t prefix, std::size_t topK) {
std::vector<std::pair<std::size_t, std::size_t>> ranked; // (hamming, id)
for (std::size_t id = 0; id < vectors.size(); ++id) {
std::size_t hamming = 0;
for (std::size_t i = 0; i < prefix; ++i) {
hamming += ((vectors[id][i] >= 0.0F) != (query[i] >= 0.0F)) ? 1U : 0U;
}
ranked.emplace_back(hamming, id);
}
std::ranges::sort(ranked);
ranked.resize(std::min(topK, ranked.size()));
return ranked;
}
} // namespace
TEST_CASE("BinaryQuantizedIndex without rotation keeps raw-coordinate sign bits",
"[vector][bq][index][characterization][catch2]") {
constexpr std::size_t kVectors = 48;
constexpr std::size_t kDimension = 96;
SplitMixFloats rng{42};
std::vector<std::size_t> ids(kVectors);
std::vector<std::vector<float>> vectors(kVectors, std::vector<float>(kDimension));
for (std::size_t id = 0; id < kVectors; ++id) {
ids[id] = id;
for (auto& value : vectors[id]) {
value = rng.next();
}
}
std::vector<float> query(kDimension);
for (auto& value : query) {
value = rng.next();
}
for (const std::size_t prefix : {std::size_t{0}, std::size_t{16}}) {
auto index = BinaryQuantizedIndex::build(ids, vectors, prefix);
REQUIRE(index.has_value());
const auto effective = prefix == 0 ? kDimension : prefix;
CHECK(index.value()->dimension() == effective);
const auto hits = index.value()->search(query, 10);
const auto expected = rawSignShortlist(vectors, query, effective, 10);
REQUIRE(hits.size() == expected.size());
for (std::size_t i = 0; i < hits.size(); ++i) {
CHECK(hits[i].hammingDistance == expected[i].first);
CHECK(hits[i].id == expected[i].second);
CHECK(hits[i].estimatedSimilarity ==
Catch::Approx(
std::cos(std::numbers::pi_v<float> * static_cast<float>(expected[i].first) /
static_cast<float>(effective))));
}
}
}
namespace {
double dotProduct(std::span<const float> a, std::span<const float> b) {
double acc = 0.0;
for (std::size_t i = 0; i < a.size(); ++i) {
acc += static_cast<double>(a[i]) * static_cast<double>(b[i]);
}
return acc;
}
double cosine(std::span<const float> a, std::span<const float> b) {
return dotProduct(a, b) / std::sqrt(dotProduct(a, a) * dotProduct(b, b));
}
// Anisotropic synthetic embeddings: three dominant coordinates carry the geometry and the
// remaining coordinates hold small noise, the regime where raw sign bits mostly encode noise.
std::vector<std::vector<float>> anisotropicVectors(std::size_t count, std::size_t dimension,
std::uint64_t seed) {
SplitMixFloats rng{seed};
std::vector<std::vector<float>> vectors(count, std::vector<float>(dimension));
for (auto& vector : vectors) {
for (auto& value : vector) {
value = 0.03F * rng.next();
}
for (std::size_t d = 0; d < 3; ++d) {
vector[d] = rng.next();
}
}
return vectors;
}
double meanCosineEstimateError(const std::vector<std::vector<float>>& vectors,
yams::vector::BinaryRotation rotation) {
std::vector<std::size_t> ids(vectors.size());
for (std::size_t i = 0; i < ids.size(); ++i) {
ids[i] = i;
}
auto index = BinaryQuantizedIndex::build(ids, vectors, 0, rotation);
REQUIRE(index.has_value());
double error = 0.0;
std::size_t pairs = 0;
for (std::size_t q = 0; q < vectors.size(); ++q) {
for (const auto& hit : index.value()->search(vectors[q], vectors.size())) {
if (hit.id == q) {
continue;
}
error += std::abs(static_cast<double>(hit.estimatedSimilarity) -
cosine(vectors[q], vectors[hit.id]));
++pairs;
}
}
return error / static_cast<double>(pairs);
}
} // namespace
TEST_CASE("Randomized Hadamard rotation is orthogonal", "[vector][bq][rotation][catch2]") {
using yams::vector::RandomizedHadamardRotation;
constexpr std::size_t kDimension = 300; // padded to 512
const RandomizedHadamardRotation rotation(kDimension, yams::vector::kDefaultBinaryRotationSeed);
CHECK(rotation.inputDimension() == kDimension);
CHECK(rotation.outputDimension() == 512U);
CHECK(rotation.seed() == yams::vector::kDefaultBinaryRotationSeed);
SplitMixFloats rng{7};
std::vector<std::vector<float>> inputs(6, std::vector<float>(kDimension));
std::vector<std::vector<float>> outputs;
for (auto& input : inputs) {
for (auto& value : input) {
value = rng.next();
}
outputs.push_back(rotation.apply(input));
REQUIRE(outputs.back().size() == 512U);
}
for (std::size_t i = 0; i < inputs.size(); ++i) {
CHECK(std::sqrt(dotProduct(outputs[i], outputs[i])) ==
Catch::Approx(std::sqrt(dotProduct(inputs[i], inputs[i]))).epsilon(1e-5));
for (std::size_t j = i + 1; j < inputs.size(); ++j) {
CHECK(dotProduct(outputs[i], outputs[j]) ==
Catch::Approx(dotProduct(inputs[i], inputs[j])).margin(1e-3));
}
}
CHECK(rotation.apply(std::vector<float>(kDimension + 1, 1.0F)).empty());
}
TEST_CASE("Rotated sign bits estimate cosine better on anisotropic vectors",
"[vector][bq][rotation][catch2]") {
const auto vectors = anisotropicVectors(64, 128, 11);
const double rawError = meanCosineEstimateError(vectors, yams::vector::BinaryRotation::None);
const double rotatedError =
meanCosineEstimateError(vectors, yams::vector::BinaryRotation::Fwht);
INFO("raw=" << rawError << " rotated=" << rotatedError);
CHECK(rotatedError < rawError);
CHECK(rotatedError < 0.5 * rawError);
}
TEST_CASE("Rotated BQ index applies its own seed to stored vectors and queries",
"[vector][bq][rotation][catch2]") {
using yams::vector::BinaryRotation;
using yams::vector::RandomizedHadamardRotation;
const auto vectors = anisotropicVectors(16, 40, 3);
std::vector<std::size_t> ids(vectors.size());
for (std::size_t i = 0; i < ids.size(); ++i) {
ids[i] = 100 + i;
}
const auto query = vectors[5];
for (const std::size_t prefix : {std::size_t{0}, std::size_t{16}}) {
auto index = BinaryQuantizedIndex::build(ids, vectors, prefix, BinaryRotation::Fwht);
REQUIRE(index.has_value());
const auto& bq = *index.value();
CHECK(bq.rotation() == BinaryRotation::Fwht);
CHECK(bq.rotationSeed() == yams::vector::kDefaultBinaryRotationSeed);
// Rotate first (40 -> 64 coordinates), then take the prefix.
CHECK(bq.dimension() == (prefix == 0 ? 64U : prefix));
const RandomizedHadamardRotation rotation(40, yams::vector::kDefaultBinaryRotationSeed);
const auto rotatedQuery = rotation.apply(query);
const auto manual =
bq.search(BinaryQuantizer::quantize(
std::span<const float>(rotatedQuery).subspan(0, bq.dimension())),
4);
const auto automatic = bq.search(query, 4);
REQUIRE(automatic.size() == manual.size());
for (std::size_t i = 0; i < manual.size(); ++i) {
CHECK(automatic[i].id == manual[i].id);
CHECK(automatic[i].hammingDistance == manual[i].hammingDistance);
}
REQUIRE_FALSE(automatic.empty());
CHECK(automatic.front().id == 105U);
CHECK(automatic.front().hammingDistance == 0U);
// A query of another dimension cannot share the rotation.
CHECK(bq.search(std::vector<float>(41, 1.0F), 4).empty());
}
auto unrotated = BinaryQuantizedIndex::build(ids, vectors);
REQUIRE(unrotated.has_value());
CHECK(unrotated.value()->rotation() == BinaryRotation::None);
CHECK(unrotated.value()->rotationSeed() == 0U);
}
Reference in a new issue View git blame Copy permalink