feat(ann): optional int8 storage for the index's vector copy
The HNSW index keeps its own copy of every vector, which at 100K x 384 f32 is ~146 MiB — the largest single item in the 2.43x footprint now that the agent stores embeddings once. `Storage::Int8` cuts that copy to a quarter by scaling each row to i8. The scale is per row, not global. Unit-length rows in d dimensions have components around 1/sqrt(d), so a fixed [-1, 1] scale spends fewer than 12 of the 255 levels on a 128-dimensional vector; measured against an exact ranking that gives 0.35 top-10 overlap. Scaling each row by its own largest component uses the full range and brings it to 0.99. Quantised distances still cost recall on their own, and `ef` does not buy it back because the loss is in the distances rather than the graph: at N=100K recall@10 tops out at 0.967 against f32's 0.9995. Re-scoring a wider candidate pool against the exact vectors removes the gap (0.9940 vs 0.9945 at ef=64) for ~13% of query throughput and ~16% of build time. That is the intended use, so it is what the test asserts — against ground truth, not against the f32 index, whose own mistakes a re-scored search is entitled to get right. Default is unchanged: `Storage::Float32`, chosen by every existing constructor. Serialized indexes carry f32 vectors and no storage tag, so a quantised index is rebuilt rather than loaded; `compact()` keeps the storage it was given. The harness grows `--int8` and `--rerank` axes, and reports the storage in each table header. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
@@ -24,7 +24,7 @@
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
|
||||
|
||||
const DIM: usize = 384;
|
||||
const K: usize = 10;
|
||||
@@ -84,6 +84,22 @@ struct Dataset {
|
||||
/// that appears only on clustered data points at graph connectivity.
|
||||
static UNIFORM: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||
|
||||
/// `--int8`: build the HNSW index over int8-quantised vectors (a quarter of
|
||||
/// the memory) instead of f32, to price the recall it costs.
|
||||
static INT8: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||
|
||||
/// `--rerank`: re-score the candidate pool against the exact vectors before
|
||||
/// taking the top K.
|
||||
static RERANK: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||
|
||||
fn storage() -> Storage {
|
||||
if INT8.load(std::sync::atomic::Ordering::Relaxed) {
|
||||
Storage::Int8
|
||||
} else {
|
||||
Storage::Float32
|
||||
}
|
||||
}
|
||||
|
||||
fn make_dataset(n: usize, seed: u64) -> Dataset {
|
||||
let mut rng = Rng(seed);
|
||||
if UNIFORM.load(std::sync::atomic::Ordering::Relaxed) {
|
||||
@@ -169,6 +185,11 @@ fn text_for(cluster: usize, i: usize, rng: &mut Rng) -> String {
|
||||
// Measurement helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Exact cosine distance between unit-length vectors.
|
||||
fn exact_dist(a: &[f32], b: &[f32]) -> f32 {
|
||||
1.0 - a.iter().zip(b).map(|(x, y)| x * y).sum::<f32>()
|
||||
}
|
||||
|
||||
fn exact_top_k(vectors: &[Vec<f32>], query: &[f32], k: usize) -> Vec<usize> {
|
||||
// Vectors are unit length, so cosine order == dot-product order.
|
||||
let mut scored: Vec<(usize, f32)> = vectors
|
||||
@@ -270,11 +291,12 @@ fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
|
||||
.collect();
|
||||
|
||||
let started = Instant::now();
|
||||
let index = HnswIndex::build_with_metric(
|
||||
let index = HnswIndex::build_with(
|
||||
&data.vectors,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
DistanceMetric::Cosine,
|
||||
storage(),
|
||||
);
|
||||
let build = started.elapsed();
|
||||
|
||||
@@ -291,7 +313,8 @@ fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
|
||||
);
|
||||
|
||||
println!(
|
||||
"\n### HNSW, N = {n}, dim = {DIM}, M = {HNSW_M}, ef_construction = {HNSW_EF_CONSTRUCTION}\n"
|
||||
"\n### HNSW, N = {n}, dim = {DIM}, M = {HNSW_M}, ef_construction = {HNSW_EF_CONSTRUCTION}, storage = {:?}\n",
|
||||
index.storage()
|
||||
);
|
||||
println!(
|
||||
"build: {:.1} ms ({:.0} vectors/s) · exact scan: {:.0} QPS, p50 {:.0} µs\n",
|
||||
@@ -302,12 +325,26 @@ fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
|
||||
);
|
||||
println!("| ef | recall@{K} | QPS | p50 µs | p99 µs |");
|
||||
println!("|---:|---:|---:|---:|---:|");
|
||||
// With a quantised index the distances it returns are approximate, so
|
||||
// the candidates are re-scored against the exact vectors the caller
|
||||
// already holds (in the agent, the embedding cache) before taking the
|
||||
// top K. `--rerank` prices that: it costs one exact distance per
|
||||
// candidate and is what decides whether int8 is usable.
|
||||
let rerank = RERANK.load(std::sync::atomic::Ordering::Relaxed);
|
||||
let pool = if rerank { K * 4 } else { K };
|
||||
for ef in EF_VALUES {
|
||||
let mut hits = 0usize;
|
||||
let mut samples = Vec::with_capacity(data.queries.len());
|
||||
for (q, want) in data.queries.iter().zip(&truth) {
|
||||
let t = Instant::now();
|
||||
let got = index.search(q, K, ef);
|
||||
let mut got = index.search(q, pool, ef.max(pool));
|
||||
if rerank {
|
||||
for cand in &mut got {
|
||||
cand.1 = exact_dist(&data.vectors[cand.0], q);
|
||||
}
|
||||
got.select_nth_unstable_by(K - 1, |a, b| a.1.total_cmp(&b.1));
|
||||
got.truncate(K);
|
||||
}
|
||||
samples.push(t.elapsed());
|
||||
hits += got.iter().filter(|(id, _)| want.contains(id)).count();
|
||||
}
|
||||
@@ -445,11 +482,12 @@ fn fusion_study(n: usize) {
|
||||
.map(|(i, c)| text_for(*c, i, &mut rng))
|
||||
.collect();
|
||||
let bm25 = BM25Index::build(&texts, &vec![0u8; n]);
|
||||
let index = HnswIndex::build_with_metric(
|
||||
let index = HnswIndex::build_with(
|
||||
&data.vectors,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
DistanceMetric::Cosine,
|
||||
storage(),
|
||||
);
|
||||
|
||||
let vec_pool = (K * 8).max(64);
|
||||
@@ -571,6 +609,14 @@ fn main() {
|
||||
}
|
||||
return;
|
||||
}
|
||||
if args.iter().any(|a| a == "--int8") {
|
||||
INT8.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
println!("(int8-quantised index vectors)");
|
||||
}
|
||||
if args.iter().any(|a| a == "--rerank") {
|
||||
RERANK.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
println!("(candidates re-scored against exact vectors)");
|
||||
}
|
||||
if args.iter().any(|a| a == "--uniform") {
|
||||
UNIFORM.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
println!("(uniform random data)");
|
||||
|
||||
Reference in New Issue
Block a user