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:
@@ -57,6 +57,52 @@ What remains at 2.43x: the flat vectors (1.0x), the HNSW index's own copy of
|
||||
them (1.0x), and text, ids and graph (~0.4x). The index copy is the next
|
||||
target — it is what a quantised or borrowed representation would address.
|
||||
|
||||
### Quantising the index copy (`quantized_index`)
|
||||
|
||||
`MemoryConfig::quantized_index` stores the index's copy as `i8` instead of
|
||||
`f32`. Same harness, same binary, `--footprint --full` with and without
|
||||
`--int8`:
|
||||
|
||||
| N | vectors (raw) | indexes, f32 | indexes, int8 | reopened, f32 | reopened, int8 |
|
||||
|---:|---:|---:|---:|---:|---:|
|
||||
| 1 000 | 1 MiB | 2 MiB | 1 MiB | 4 MiB (2.40x) | 2 MiB (1.64x) |
|
||||
| 10 000 | 15 MiB | 32 MiB | 14 MiB | 44 MiB (3.03x) | 27 MiB (1.81x) |
|
||||
| 100 000 | 146 MiB | 266 MiB | **123 MiB** | 399 MiB (2.72x) | **256 MiB (1.74x)** |
|
||||
|
||||
The scale is **per row**, not global. A unit-length row in `d` dimensions has
|
||||
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 — unusable. Scaling each row by its own
|
||||
largest component brings the same measurement 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 in the graph
|
||||
(`--ann-only --full`, N = 100 000):
|
||||
|
||||
| ef | recall@10, f32 | recall@10, int8 | recall@10, int8 + re-score |
|
||||
|---:|---:|---:|---:|
|
||||
| 32 | 0.9775 | 0.9415 | 0.9785 |
|
||||
| 64 | 0.9945 | 0.9625 | 0.9940 |
|
||||
| 128 | 0.9995 | 0.9670 | 0.9990 |
|
||||
| 256 | 0.9995 | 0.9670 (ceiling) | 0.9990 |
|
||||
|
||||
Re-scoring closes the gap: the store already holds the exact embeddings, so
|
||||
the query path re-scores the candidate pool against them before fusion. That
|
||||
is done automatically whenever the index is quantised. What it costs is
|
||||
throughput — about 13% of QPS and 16% of build time at 100 000 x 384. So the
|
||||
setting trades ~13% of query speed for ~36% of the process's memory at equal
|
||||
recall. It is **off by default**: the right side of that trade depends on
|
||||
whether the deployment is short of memory or short of CPU.
|
||||
|
||||
A measurement trap worth recording: the synthetic `clustered` generator in the
|
||||
`clawhdf5-ann` tests draws clusters far tighter than any real embedding, so
|
||||
neighbours there sit closer together than the quantisation error and top-10
|
||||
*identity* is noise. Scored on that fixture int8 looks catastrophic (0.57
|
||||
overlap) — a fact about the fixture, not the storage. The tests use random
|
||||
vectors, and recall is measured against brute-force ground truth rather than
|
||||
against the f32 index, whose own approximation errors a re-scored search is
|
||||
entitled to get right.
|
||||
|
||||
## Read harness
|
||||
|
||||
Produced by `cargo run --release -p clawhdf5-bench --bin read_harness`: a 4096 x
|
||||
|
||||
@@ -2,6 +2,26 @@
|
||||
|
||||
## Unreleased
|
||||
|
||||
### Memory
|
||||
- `clawhdf5-agent`: **`MemoryConfig::quantized_index`** stores the vector
|
||||
index's own copy of the embeddings as `i8` rather than `f32`, which at 100k
|
||||
384-dim entries takes the index from 266 to 123 MiB and the whole reopened
|
||||
store from 399 to 256 MiB (2.72x -> **1.74x** the raw vectors). Quantised
|
||||
distances are approximate and `ef` cannot compensate — recall@10 tops out at
|
||||
0.967 against f32's 0.9995 — so the query path re-scores the candidate pool
|
||||
against the exact embeddings the store already holds, which restores recall
|
||||
(0.9940 vs 0.9945 at ef=64) for about 13% of QPS. **Off by default**: it
|
||||
trades query speed for memory, and which side is worth more depends on the
|
||||
deployment. The setting is persisted, so a reopened store does not silently
|
||||
revert to four times the index memory.
|
||||
- `clawhdf5-ann`: `Storage::Int8` and the `build_with` / `new_with` /
|
||||
`from_graph_bytes_with` constructors that select it. The scale is per row,
|
||||
not global — a fixed `[-1, 1]` scale spends fewer than 12 of the 255 levels
|
||||
on a unit-length 128-dim vector and is unusable (0.35 top-10 overlap against
|
||||
an exact ranking, versus 0.99 per row). `compact()` keeps the storage it was
|
||||
given; serialized indexes still carry f32 vectors, so a quantised index is
|
||||
rebuilt rather than loaded.
|
||||
|
||||
### Memory
|
||||
- `clawhdf5-agent`: **a loaded store holds ~30% less memory** (100k 384-dim
|
||||
entries: 505 -> 357 MiB, 3.44x -> 2.43x the raw vectors). The cache kept
|
||||
|
||||
@@ -39,7 +39,13 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F
|
||||
(plain closest-M capped recall on clustered data: 0.31 recall@10 at 100K). Its
|
||||
graph is saved to `<store>.h5.ann` at each checkpoint and reloaded by `open()`
|
||||
(tied to the checkpoint by a generation id; stale/damaged sidecars are
|
||||
ignored and the index rebuilt). `hybrid_search` keeps one incremental BM25
|
||||
ignored and the index rebuilt). `MemoryConfig::quantized_index` (off by
|
||||
default, persisted) stores the index's own copy of the embeddings as `i8`,
|
||||
which roughly halves a loaded store's memory (2.72x -> 1.74x the raw vectors
|
||||
at 100K); because quantised distances are approximate and `ef` cannot
|
||||
compensate, the query path then re-scores the candidate pool against the
|
||||
exact embeddings, which holds recall at the f32 index's level and costs
|
||||
~13% of QPS. `hybrid_search` keeps one incremental BM25
|
||||
index for the life of the store and never writes the store: Hebbian
|
||||
activation boosts are persisted by the next checkpoint (or on drop), not per
|
||||
query. Measure any search-path change with
|
||||
|
||||
@@ -432,6 +432,13 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
|
||||
| `agent` | no | Full agent memory layer |
|
||||
| `float16` | **yes** | Half-precision embedding storage (2× compression) |
|
||||
| `hnsw` | **yes** | HNSW approximate vector index for `hybrid_search` (via `clawhdf5-ann`); disable for an exact linear scan |
|
||||
|
||||
`MemoryConfig::quantized_index` (off by default) stores the HNSW index's own
|
||||
copy of the embeddings as `i8`, roughly halving a loaded store's memory
|
||||
(2.72x -> 1.74x the raw vectors at 100k x 384). Quantised distances are
|
||||
approximate, so the query path re-scores the candidate pool against the exact
|
||||
embeddings the store already holds — recall matches the `f32` index, at about
|
||||
13% fewer queries per second. See `BENCHMARKS.md`, "Quantising the index copy".
|
||||
| `parallel` | no | Rayon parallel search |
|
||||
| `fast-math` | no | BLAS matrix-vector multiply |
|
||||
| `accelerate` | no | Apple Accelerate / AMX (macOS) |
|
||||
|
||||
@@ -118,6 +118,7 @@ mod tests {
|
||||
created_at: "2025-01-01T00:00:00Z".to_string(),
|
||||
wal_enabled: false,
|
||||
wal_max_entries: 500,
|
||||
quantized_index: false,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -62,7 +62,7 @@ use std::path::{Path, PathBuf};
|
||||
|
||||
use cache::MemoryCache;
|
||||
#[cfg(feature = "hnsw")]
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
|
||||
use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||
|
||||
/// HNSW construction parameters used for the agent's vector index. Cosine is the
|
||||
@@ -139,6 +139,17 @@ pub struct MemoryConfig {
|
||||
pub created_at: String,
|
||||
pub wal_enabled: bool,
|
||||
pub wal_max_entries: usize,
|
||||
/// Store the vector index's own copy of the embeddings as int8 rather than
|
||||
/// f32, a quarter of the memory.
|
||||
///
|
||||
/// The index's copy is the single largest part of a loaded store's
|
||||
/// footprint. Quantised distances are approximate, so the candidate pool
|
||||
/// is re-scored against the cache's exact embeddings before fusion, which
|
||||
/// restores recall; what it costs is throughput — roughly 13% of queries
|
||||
/// per second and 16% of build time at 100K x 384. See `BENCHMARKS.md`.
|
||||
///
|
||||
/// Has no effect without the `hnsw` feature.
|
||||
pub quantized_index: bool,
|
||||
}
|
||||
|
||||
impl MemoryConfig {
|
||||
@@ -160,6 +171,7 @@ impl MemoryConfig {
|
||||
created_at,
|
||||
wal_enabled: true,
|
||||
wal_max_entries: 500,
|
||||
quantized_index: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -444,7 +456,17 @@ impl HDF5Memory {
|
||||
|
||||
#[cfg(feature = "hnsw")]
|
||||
let loaded_index = if replay_only_appended {
|
||||
Self::load_vector_index(path, checkpoint.ann_generation, &cache, n_checkpoint)
|
||||
Self::load_vector_index(
|
||||
path,
|
||||
checkpoint.ann_generation,
|
||||
&cache,
|
||||
n_checkpoint,
|
||||
if config.quantized_index {
|
||||
Storage::Int8
|
||||
} else {
|
||||
Storage::Float32
|
||||
},
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -558,6 +580,7 @@ impl HDF5Memory {
|
||||
generation: Option<u64>,
|
||||
cache: &MemoryCache,
|
||||
n_checkpoint: usize,
|
||||
storage: Storage,
|
||||
) -> Option<HnswIndex> {
|
||||
let generation = generation?;
|
||||
let bytes = std::fs::read(Self::vector_index_path(store)).ok()?;
|
||||
@@ -568,7 +591,7 @@ impl HDF5Memory {
|
||||
let vectors: Vec<Vec<f32>> = (0..n_checkpoint)
|
||||
.map(|i| cache.embeddings.get(i).map(<[f32]>::to_vec))
|
||||
.collect::<Option<_>>()?;
|
||||
let mut index = HnswIndex::from_graph_bytes(graph, vectors).ok()?;
|
||||
let mut index = HnswIndex::from_graph_bytes_with(graph, vectors, storage).ok()?;
|
||||
if index.dimension() != cache.embedding_dim {
|
||||
return None;
|
||||
}
|
||||
@@ -803,6 +826,16 @@ impl HDF5Memory {
|
||||
// the index length drifts from the cache length (covering any mutation path
|
||||
// that doesn't call a hook, e.g. consolidation pushes).
|
||||
|
||||
/// How the index should store its copy of the vectors, per the config.
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn index_storage(&self) -> Storage {
|
||||
if self.config.quantized_index {
|
||||
Storage::Int8
|
||||
} else {
|
||||
Storage::Float32
|
||||
}
|
||||
}
|
||||
|
||||
/// Build an HNSW index over the entire cache, re-applying tombstones as
|
||||
/// soft-deletions so node ids stay aligned with cache indices.
|
||||
///
|
||||
@@ -821,11 +854,12 @@ impl HDF5Memory {
|
||||
// The index owns its vectors, so it needs rows rather than the cache's
|
||||
// flat buffer. This copy is the index's own; the cache keeps one.
|
||||
let rows: Vec<Vec<f32>> = self.cache.embeddings.iter().map(<[f32]>::to_vec).collect();
|
||||
let mut index = HnswIndex::build_with_metric(
|
||||
let mut index = HnswIndex::build_with(
|
||||
&rows,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
DistanceMetric::Cosine,
|
||||
self.index_storage(),
|
||||
);
|
||||
for (i, &t) in self.cache.tombstones.iter().enumerate() {
|
||||
if t != 0 {
|
||||
|
||||
@@ -104,6 +104,10 @@ pub fn build_hdf5_file_with_meta(
|
||||
"wal_max_entries",
|
||||
AttrValue::I64(config.wal_max_entries as i64),
|
||||
);
|
||||
meta.set_attr(
|
||||
"quantized_index",
|
||||
AttrValue::I64(config.quantized_index.into()),
|
||||
);
|
||||
meta.set_attr(
|
||||
"edgehdf5_version",
|
||||
AttrValue::String(ZEROCLAW_VERSION.into()),
|
||||
@@ -484,6 +488,7 @@ pub fn validate_and_load(
|
||||
wal_max_entries: optional_i64_attr(&attrs, "wal_max_entries")
|
||||
.and_then(|v| usize::try_from(v).ok())
|
||||
.unwrap_or(500),
|
||||
quantized_index: optional_bool_attr(&attrs, "quantized_index", false),
|
||||
};
|
||||
|
||||
// Load /memory group
|
||||
|
||||
@@ -29,10 +29,26 @@ impl HDF5Memory {
|
||||
// Over-fetch so the merge sees a useful vector pool; cosine
|
||||
// distance from the index converts back to similarity (1 - d).
|
||||
let pool = (k * 8).max(64);
|
||||
let vec_scores: Vec<(usize, f32)> = index
|
||||
.search(query_embedding, pool, pool)
|
||||
let candidates = index.search(query_embedding, pool, pool);
|
||||
// A quantised index returns approximate distances, and no
|
||||
// amount of `ef` fixes that — the loss is in the distances,
|
||||
// not the graph. Re-score the pool against the cache's exact
|
||||
// embeddings, which cost nothing extra to keep: recall then
|
||||
// matches an f32 index. See `BENCHMARKS.md`.
|
||||
let exact = index.storage() == clawhdf5_ann::Storage::Int8;
|
||||
let vec_scores: Vec<(usize, f32)> = candidates
|
||||
.into_iter()
|
||||
.map(|(id, dist)| (id, 1.0 - dist))
|
||||
.map(|(id, dist)| {
|
||||
let score = if exact {
|
||||
crate::vector_search::cosine_similarity(
|
||||
query_embedding,
|
||||
&self.cache.embeddings[id],
|
||||
)
|
||||
} else {
|
||||
1.0 - dist
|
||||
};
|
||||
(id, score)
|
||||
})
|
||||
.collect();
|
||||
// Fusion normalises over every keyword match, so it needs all
|
||||
// the scores — but not ranked.
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -563,7 +563,9 @@ fn bench_footprint(n: usize) {
|
||||
.collect();
|
||||
let after_entries = heap_bytes();
|
||||
|
||||
let mut mem = HDF5Memory::create(MemoryConfig::new(path, "bench", DIM)).unwrap();
|
||||
let mut config = MemoryConfig::new(path, "bench", DIM);
|
||||
config.quantized_index = INT8.load(std::sync::atomic::Ordering::Relaxed);
|
||||
let mut mem = HDF5Memory::create(config).unwrap();
|
||||
mem.save_batch(entries).unwrap();
|
||||
let after_store = heap_bytes();
|
||||
|
||||
|
||||
@@ -28,6 +28,10 @@ enum Commands {
|
||||
/// Enable write-ahead log
|
||||
#[arg(long)]
|
||||
wal: bool,
|
||||
/// Store the vector index's copy of the embeddings as int8, roughly
|
||||
/// halving a loaded store's memory at about 13% fewer queries/second
|
||||
#[arg(long)]
|
||||
quantized_index: bool,
|
||||
},
|
||||
/// Save a memory entry (reads JSON from stdin or --json)
|
||||
Save {
|
||||
@@ -88,9 +92,15 @@ fn main() {
|
||||
|
||||
fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
||||
match cli.command {
|
||||
Commands::Create { agent_id, dim, wal } => {
|
||||
Commands::Create {
|
||||
agent_id,
|
||||
dim,
|
||||
wal,
|
||||
quantized_index,
|
||||
} => {
|
||||
let mut config = MemoryConfig::new(cli.path.clone(), &agent_id, dim);
|
||||
config.wal_enabled = wal;
|
||||
config.quantized_index = quantized_index;
|
||||
let mem = HDF5Memory::create(config)?;
|
||||
let j = serde_json::json!({
|
||||
"status": "created",
|
||||
@@ -98,6 +108,7 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
||||
"agent_id": agent_id,
|
||||
"embedding_dim": dim,
|
||||
"wal_enabled": wal,
|
||||
"quantized_index": quantized_index,
|
||||
"count": mem.count(),
|
||||
});
|
||||
println!("{}", serde_json::to_string_pretty(&j)?);
|
||||
|
||||
@@ -364,6 +364,11 @@ cargo install --path crates/clawhdf5-cli
|
||||
clawhdf5 --path agent.h5 create --agent-id my-agent --dim 384 --wal
|
||||
```
|
||||
|
||||
Add `--quantized-index` to store the vector index's copy of the embeddings as
|
||||
int8. That roughly halves a loaded store's memory at about 13% fewer queries
|
||||
per second, with recall unchanged — the query path re-scores candidates
|
||||
against the exact embeddings. The setting is recorded in the file.
|
||||
|
||||
Output:
|
||||
```json
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user