feat(agent): optional int8 vector index, re-scored against exact embeddings
`MemoryConfig::quantized_index` stores the HNSW index's own copy of the embeddings as i8 rather than f32. At 100k x 384 that takes the index from 266 to 123 MiB and the whole reopened store from 399 to 256 MiB — 2.72x to 1.74x the raw vectors, the largest remaining item in the footprint. Quantised distances are approximate and `ef` cannot compensate, because the loss is in the distances rather than in the graph: recall@10 tops out at 0.967 against f32's 0.9995 and does not move between ef=128 and ef=256. The store already holds the exact embeddings, though, so when the index is quantised the query path re-scores the candidate pool against them before fusion. That restores recall (0.9940 vs 0.9945 at ef=64) and costs about 13% of QPS. Off by default: it trades query speed for memory and which side is worth more depends on the deployment. The flag is persisted in `/meta`, so a reopened store does not silently revert to four times the index memory, and the sidecar graph is rehydrated into the configured storage. Also on the CLI as `create --quantized-index`. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
@@ -165,3 +165,72 @@ fn save_batch_then_search_is_consistent() {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quantized_index_matches_the_f32_index_after_re_scoring() {
|
||||
// A quantised index holds approximate vectors, but the store still has the
|
||||
// exact ones, so the query path re-scores the candidate pool before
|
||||
// fusion. The results a caller sees should therefore be the same.
|
||||
let dim = 64;
|
||||
let n = 400;
|
||||
let mut seed = 0x5EED_1234_5678_9ABC;
|
||||
let vectors: Vec<Vec<f32>> = (0..n).map(|_| make_vector(&mut seed, dim)).collect();
|
||||
let queries: Vec<Vec<f32>> = (0..20).map(|_| make_vector(&mut seed, dim)).collect();
|
||||
|
||||
let build = |dir: &TempDir, quantized: bool| {
|
||||
let mut config = MemoryConfig::new(dir.path().join("mem.h5"), "agent", dim);
|
||||
config.quantized_index = quantized;
|
||||
let mut mem = HDF5Memory::create(config).unwrap();
|
||||
for (i, v) in vectors.iter().enumerate() {
|
||||
mem.save(entry(&format!("chunk {i}"), v.clone(), &format!("k{i}")))
|
||||
.unwrap();
|
||||
}
|
||||
mem
|
||||
};
|
||||
|
||||
let exact_dir = TempDir::new().unwrap();
|
||||
let quant_dir = TempDir::new().unwrap();
|
||||
let mut exact = build(&exact_dir, false);
|
||||
let mut quantized = build(&quant_dir, true);
|
||||
|
||||
let k = 10;
|
||||
let mut agree = 0;
|
||||
for q in &queries {
|
||||
let want: Vec<usize> = exact
|
||||
.hybrid_search(q, "", 1.0, 0.0, k)
|
||||
.iter()
|
||||
.map(|r| r.index)
|
||||
.collect();
|
||||
agree += quantized
|
||||
.hybrid_search(q, "", 1.0, 0.0, k)
|
||||
.iter()
|
||||
.filter(|r| want.contains(&r.index))
|
||||
.count();
|
||||
}
|
||||
let overlap = agree as f64 / (k * queries.len()) as f64;
|
||||
assert!(
|
||||
overlap >= 0.95,
|
||||
"quantised store should match the f32 one: {overlap}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quantized_index_setting_survives_a_reopen() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("mem.h5");
|
||||
let mut config = MemoryConfig::new(path.clone(), "agent", 8);
|
||||
config.quantized_index = true;
|
||||
let mut mem = HDF5Memory::create(config).unwrap();
|
||||
let mut seed = 7;
|
||||
for i in 0..30 {
|
||||
mem.save(entry(&format!("c{i}"), make_vector(&mut seed, 8), "t"))
|
||||
.unwrap();
|
||||
}
|
||||
mem.flush_wal().unwrap();
|
||||
drop(mem);
|
||||
|
||||
// Reopening must not silently quadruple the index's memory, so the flag
|
||||
// is part of the stored config rather than a per-session choice.
|
||||
let reopened = HDF5Memory::open(&path).unwrap();
|
||||
assert!(reopened.config().quantized_index);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user