8aa1fbb414
Three real defects, plus a retracted fourth that was a probe artifact. P3 (fixed) - query/stored normalization asymmetry. The write path L2-normalized every stored vector; the read path passed the caller's raw query straight to the index, so the two sides lived in different spaces. With unit v, d = |q|^2 - 2q.v + 1, so a non-unit query shifted and scaled every distance by |q|^2. Measured live: 591-1174 against a documented [0,4], and an exact match scoring |q|^2 - 1 instead of ~0. vector_search_items now normalizes with the canonical l2_normalize; a zero-norm query (no direction, so nearest-by-cosine is undefined) is rejected with 400. Ranking is unchanged - |q|^2 and 1 are constant across candidates - which is why it went unnoticed; what broke was every absolute use of the number. WIRE-VISIBLE, recorded in CHANGELOG. P2 (fixed) - the blob path had zero instrumentation. Added per-kind tidaldb_cluster_blobs_originated/applied/apply_failed totals. Label cardinality is fixed at 4 by construction via a new BlobKind enum, and BlobRecord::blob_kind is now the ONE exhaustive match over the variants (kind() derives from it), so a new variant is a compile error in one place instead of a silent zero in three. Only the live apply path is counted - boot replay would inflate applied past originated on every restart. Coverage gap (fixed) - tidaldb_usearch_vector_count rendered only the metrics owner's shard group, so on a 3-group node two thirds of the corpus had no vector-count series at all. Co-located groups now render shard="N"; the owner stays unlabeled for wire compatibility, so an alert grouped by (shard) buckets each replica set separately without double-counting. P1 (RETRACTED) - the "replica-divergent vector index" does not exist. Every probe wrote through the /sharded/ surface, which hash-partitions and applies to the owning region's local store with no WAL append, and therefore does not replicate BY DESIGN (cluster/node.rs:8828-8829). A controlled A/B settled it: on /items plus /embeddings all 6 entities reach all 3 replicas; on the sharded surface four of six reach exactly one node. Both are now pinned by tests. See tmp/vector-search-correctness/diagnosis.md and the k3s-fleet cluster-state.yaml entry RETRACTED_blob_replication_rf1_2026_08_30. Pre-work: usearch_index.rs 872 to 503 lines by extracting its tests to a sibling (the project's existing path-attribute convention), and the three hand-rolled l2_normalize copies collapsed to one. The two entity copies used a zero threshold about 2900x looser than the canonical one; normalize_centroid now names the centroid zero-tolerance policy once, and a test pins the tightened behavior. Tests: 2107 lib (+5), 8 vector_search e2e (+4, three of which fail without the P3 fix), 4 cluster_sharding e2e (+2). The heavy multiproc tests in cluster_sharding are now serialized - four concurrent 3-node clusters made the pre-existing failover test miss its 10s budget.
879 lines
35 KiB
Rust
879 lines
35 KiB
Rust
//! Preference vector: per-user taste embedding with L2 normalization invariant.
|
||
//!
|
||
//! Tracks user taste by maintaining a preference vector that evolves with
|
||
//! interactions. The vector is L2-normalized on every update to ensure
|
||
//! consistent cosine similarity scoring during personalized ranking.
|
||
//!
|
||
//! # Durability
|
||
//!
|
||
//! Preference vectors are DERIVED state with no WAL backing (the signal record
|
||
//! carries no `for_user`), so they are persisted by [`PreferenceVectors::checkpoint`]
|
||
//! (called on the periodic checkpoint thread, on backup, and on shutdown) and
|
||
//! restored on open by [`PreferenceVectors::restore`]. This bounds the post-crash
|
||
//! loss window to the checkpoint interval rather than losing every update since
|
||
//! the last clean shutdown.
|
||
|
||
use dashmap::DashMap;
|
||
|
||
use crate::storage::vector::{VectorError, l2_normalize_in_place};
|
||
|
||
/// Per-user preference vector, L2-normalized.
|
||
///
|
||
/// The vector is updated via exponential moving average: each new interaction
|
||
/// embedding is blended with the current preference using a learning rate.
|
||
///
|
||
/// Thread-safe via `DashMap` -- concurrent updates to different users never
|
||
/// contend.
|
||
pub struct PreferenceVectors {
|
||
/// `user_id` -> normalized preference vector
|
||
inner: DashMap<u64, Vec<f32>>,
|
||
/// Dimensionality of the embedding space. All vectors must have this length.
|
||
dim: usize,
|
||
/// Base learning rate for exponential moving average updates.
|
||
/// Default: 0.1. The adaptive rate decays as `base / (1 + ln(count + 1))`.
|
||
learning_rate: f32,
|
||
/// Per-user update counts for adaptive learning rate computation.
|
||
///
|
||
/// # Known Limitation
|
||
///
|
||
/// `update_counts` is in-memory only and is NOT persisted to storage. On
|
||
/// process restart, all users start at count=0 and their adaptive learning
|
||
/// rate resets to `base_alpha`. This means the learning rate decays more
|
||
/// slowly across session boundaries than within a single session.
|
||
/// Persisting update counts requires a storage write on every preference
|
||
/// update, which introduces non-trivial overhead on the hot path; this
|
||
/// trade-off is deferred to M7 (Production Hardening).
|
||
update_counts: DashMap<u64, u64>,
|
||
}
|
||
|
||
impl PreferenceVectors {
|
||
/// Create a new preference vector store for the given embedding dimensionality.
|
||
#[must_use]
|
||
pub fn new(dim: usize) -> Self {
|
||
Self {
|
||
inner: DashMap::new(),
|
||
dim,
|
||
learning_rate: 0.1,
|
||
update_counts: DashMap::new(),
|
||
}
|
||
}
|
||
|
||
/// Create with a custom base learning rate.
|
||
#[must_use]
|
||
pub fn with_learning_rate(dim: usize, learning_rate: f32) -> Self {
|
||
Self {
|
||
inner: DashMap::new(),
|
||
dim,
|
||
learning_rate,
|
||
update_counts: DashMap::new(),
|
||
}
|
||
}
|
||
|
||
/// Get the current preference vector for a user (cloned).
|
||
///
|
||
/// Returns `None` if no preference has been recorded.
|
||
#[must_use]
|
||
pub fn get(&self, user_id: u64) -> Option<Vec<f32>> {
|
||
self.inner.get(&user_id).map(|r| r.clone())
|
||
}
|
||
|
||
/// Set the preference vector directly (e.g., from cold-start initialization).
|
||
///
|
||
/// The vector is L2-normalized before storage. Returns `false` if the
|
||
/// dimension does not match.
|
||
#[must_use]
|
||
pub fn set(&self, user_id: u64, mut vec: Vec<f32>) -> bool {
|
||
if vec.len() != self.dim {
|
||
return false;
|
||
}
|
||
normalize_centroid(&mut vec);
|
||
self.inner.insert(user_id, vec);
|
||
true
|
||
}
|
||
|
||
/// Update a user's preference vector by blending with an interaction embedding.
|
||
///
|
||
/// Uses exponential moving average with an adaptive learning rate:
|
||
/// `alpha = base_alpha / (1 + ln(update_count + 1))`
|
||
/// `pref = (1 - alpha) * pref + alpha * interaction`
|
||
/// then L2-normalizes the result.
|
||
///
|
||
/// The adaptive rate decays logarithmically with the number of updates,
|
||
/// so early interactions have high influence and later interactions refine
|
||
/// the preference more gently.
|
||
///
|
||
/// If no preference exists yet, the interaction embedding becomes the initial
|
||
/// preference (after normalization). Uses `Entry::Occupied`/`Entry::Vacant`
|
||
/// to avoid double-applying the blend on first insertion.
|
||
///
|
||
/// Returns `false` if the interaction embedding dimension does not match.
|
||
#[must_use]
|
||
pub fn update(&self, user_id: u64, interaction_embedding: &[f32]) -> bool {
|
||
if interaction_embedding.len() != self.dim {
|
||
return false;
|
||
}
|
||
|
||
// Compute adaptive learning rate: alpha = base / (1 + ln(count + 1)).
|
||
let count = {
|
||
let mut entry = self.update_counts.entry(user_id).or_insert(0);
|
||
let c = *entry;
|
||
*entry += 1;
|
||
c
|
||
};
|
||
#[allow(clippy::cast_precision_loss)]
|
||
let lr = (f64::from(self.learning_rate) / (1.0 + (count as f64).ln_1p())) as f32;
|
||
|
||
self.blend_and_normalize(user_id, interaction_embedding, lr);
|
||
true
|
||
}
|
||
|
||
/// Update a user's preference vector by blending with an interaction embedding
|
||
/// using an explicit learning rate (overrides the stored rate).
|
||
///
|
||
/// Same EMA formula as [`update`](Self::update) but with `lr` parameter:
|
||
/// `pref = (1 - lr) * pref + lr * interaction`
|
||
///
|
||
/// Returns `false` if the interaction embedding dimension does not match.
|
||
#[must_use]
|
||
pub fn update_with_custom_rate(
|
||
&self,
|
||
user_id: u64,
|
||
interaction_embedding: &[f32],
|
||
lr: f32,
|
||
) -> bool {
|
||
if interaction_embedding.len() != self.dim {
|
||
return false;
|
||
}
|
||
|
||
self.blend_and_normalize(user_id, interaction_embedding, lr);
|
||
true
|
||
}
|
||
|
||
/// Blend `interaction_embedding` into the user's preference at learning
|
||
/// rate `lr` and re-normalize, the shared EMA body for both [`update`] and
|
||
/// [`update_with_custom_rate`].
|
||
///
|
||
/// An existing preference is blended in place via
|
||
/// `pref = (1 - lr) * pref + lr * interaction` and re-normalized; if no
|
||
/// preference exists yet, the (normalized) interaction embedding becomes the
|
||
/// initial preference, avoiding a double-applied blend on first insertion.
|
||
///
|
||
/// The caller is responsible for the dimension check — `dim` mismatch is
|
||
/// validated before this is reached, and the per-element `zip` would
|
||
/// otherwise silently blend only the overlapping prefix.
|
||
///
|
||
/// [`update`]: Self::update
|
||
/// [`update_with_custom_rate`]: Self::update_with_custom_rate
|
||
fn blend_and_normalize(&self, user_id: u64, interaction_embedding: &[f32], lr: f32) {
|
||
use dashmap::mapref::entry::Entry;
|
||
|
||
match self.inner.entry(user_id) {
|
||
Entry::Occupied(mut occ) => {
|
||
let pref = occ.get_mut();
|
||
for (p, &i) in pref.iter_mut().zip(interaction_embedding.iter()) {
|
||
*p = (1.0 - lr).mul_add(*p, lr * i);
|
||
}
|
||
normalize_centroid(pref);
|
||
}
|
||
Entry::Vacant(vac) => {
|
||
let mut v = interaction_embedding.to_vec();
|
||
normalize_centroid(&mut v);
|
||
vac.insert(v);
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Compute cosine similarity between a user's preference and a candidate embedding.
|
||
///
|
||
/// Returns `None` if the user has no preference vector or dimensions mismatch.
|
||
/// The stored preference is L2-normalized; the candidate is normalized on-the-fly
|
||
/// so callers do not need to pre-normalize.
|
||
#[must_use]
|
||
#[allow(clippy::significant_drop_tightening)]
|
||
pub fn cosine_similarity(&self, user_id: u64, candidate: &[f32]) -> Option<f32> {
|
||
if candidate.len() != self.dim {
|
||
return None;
|
||
}
|
||
let pref = self.inner.get(&user_id)?;
|
||
let dot: f32 = pref.iter().zip(candidate.iter()).map(|(a, b)| a * b).sum();
|
||
// Divide by the candidate's L2 norm to get true cosine similarity.
|
||
// The stored preference is already unit-length, so we only need
|
||
// to normalize the candidate side.
|
||
let candidate_norm: f32 = candidate.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||
if candidate_norm < f32::EPSILON {
|
||
return Some(0.0);
|
||
}
|
||
Some(dot / candidate_norm)
|
||
}
|
||
|
||
/// Remove every stored preference vector and update count.
|
||
///
|
||
/// Used to return the store to its just-opened state so a full rebuild
|
||
/// (re-seed + signal replay) does not fold on top of already-folded vectors
|
||
/// — i.e. to keep an out-of-band reindex idempotent with a fresh boot.
|
||
pub fn clear(&self) {
|
||
self.inner.clear();
|
||
self.update_counts.clear();
|
||
}
|
||
|
||
/// Remove a single user's cold-start vector and update count.
|
||
///
|
||
/// Used by [`MultiPreferenceVectors`](crate::entities::MultiPreferenceVectors)
|
||
/// when a user crosses the warm threshold: their cold-start vector is migrated
|
||
/// into cluster 0 and then dropped here so a checkpoint never double-stores a
|
||
/// stale K=1 row alongside the cluster rows.
|
||
pub fn remove(&self, user_id: u64) {
|
||
self.inner.remove(&user_id);
|
||
self.update_counts.remove(&user_id);
|
||
}
|
||
|
||
/// Insert an already-restored (normalized) vector + update count directly,
|
||
/// bypassing the EMA blend.
|
||
///
|
||
/// The multi-vector restore path uses this to load a legacy single-vector
|
||
/// checkpoint row as a cold-start user. The caller is responsible for having
|
||
/// re-normalized the vector at the load boundary (matching [`restore`]'s
|
||
/// contract); the dimension is trusted because the caller already gated on it.
|
||
///
|
||
/// [`restore`]: Self::restore
|
||
pub fn insert_restored(&self, user_id: u64, vec: Vec<f32>, update_count: u64) {
|
||
self.inner.insert(user_id, vec);
|
||
self.update_counts.insert(user_id, update_count);
|
||
}
|
||
|
||
/// Whether this store holds a vector for `user_id`, without allocating
|
||
/// (an O(1) hash lookup). Used by
|
||
/// [`MultiPreferenceVectors::contains`](crate::entities::MultiPreferenceVectors::contains)
|
||
/// for serve-path existence guards.
|
||
#[must_use]
|
||
pub fn contains(&self, user_id: u64) -> bool {
|
||
self.inner.contains_key(&user_id)
|
||
}
|
||
|
||
/// Stage one legacy single-vector row per stored user into an existing
|
||
/// `WriteBatch` (`[count:8 LE][dim:4 LE][f32*dim]`), without performing its
|
||
/// own delete sweep. `skip(user_id) == true` omits that user's row.
|
||
///
|
||
/// [`checkpoint`](Self::checkpoint) is the standalone single-vector path (it
|
||
/// owns the delete sweep + write). This variant lets
|
||
/// [`MultiPreferenceVectors`](crate::entities::MultiPreferenceVectors)'s
|
||
/// unified checkpoint fold cold-start users into the same atomic batch as the
|
||
/// cluster rows, so a single swap covers both tiers. The `skip` predicate lets
|
||
/// that caller exclude any user already written as a warm cluster row, so a
|
||
/// legacy row can never collide with (and clobber) a warm row under the shared
|
||
/// `Tag::Preference` key.
|
||
pub fn append_legacy_rows(
|
||
&self,
|
||
batch: &mut crate::storage::WriteBatch,
|
||
skip: impl Fn(u64) -> bool,
|
||
) {
|
||
use crate::{
|
||
schema::EntityId,
|
||
storage::{Tag, encode_key},
|
||
};
|
||
for entry in &self.inner {
|
||
let user_id = *entry.key();
|
||
if skip(user_id) {
|
||
continue;
|
||
}
|
||
let value = encode_legacy_row(self.update_count(user_id), entry.value());
|
||
let key = encode_key(EntityId::new(0), Tag::Preference, &user_id.to_be_bytes());
|
||
batch.put(key, value);
|
||
}
|
||
}
|
||
|
||
/// Number of users with stored preferences.
|
||
#[must_use]
|
||
pub fn len(&self) -> usize {
|
||
self.inner.len()
|
||
}
|
||
|
||
/// Whether no preferences are stored.
|
||
#[must_use]
|
||
pub fn is_empty(&self) -> bool {
|
||
self.inner.is_empty()
|
||
}
|
||
|
||
/// Get the current update count for a user (number of times `update` has been called).
|
||
#[must_use]
|
||
pub fn update_count(&self, user_id: u64) -> u64 {
|
||
self.update_counts.get(&user_id).map_or(0, |c| *c)
|
||
}
|
||
|
||
/// Compute the adaptive learning rate for a user at their current update count.
|
||
///
|
||
/// Formula: `base_alpha / (1 + ln(count + 1))`.
|
||
#[must_use]
|
||
#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
|
||
pub fn adaptive_lr(&self, user_id: u64) -> f32 {
|
||
let count = self.update_count(user_id);
|
||
(f64::from(self.learning_rate) / (1.0 + (count as f64).ln_1p())) as f32
|
||
}
|
||
}
|
||
|
||
// ── Checkpoint / Restore ─────────────────────────────────────────────────────
|
||
|
||
impl PreferenceVectors {
|
||
/// Checkpoint every per-user preference vector + update count to durable
|
||
/// storage under `Tag::Preference`, one row per user.
|
||
///
|
||
/// All rows share the sentinel entity id (`0`) so [`restore`](Self::restore)
|
||
/// can target them with a single prefix scan. The swap is staged into one
|
||
/// `WriteBatch` and committed atomically (matching the co-engagement /
|
||
/// community checkpoints): stale rows are deleted and current rows written in
|
||
/// the same batch, so a crash mid-checkpoint never leaves a half-written set.
|
||
///
|
||
/// Key suffix: `[user: 8B BE]`. Value: `[update_count: 8B LE][dim: 4B LE][f32 * dim LE]`.
|
||
///
|
||
/// **In-engine this is superseded by
|
||
/// [`MultiPreferenceVectors`](crate::entities::MultiPreferenceVectors), which
|
||
/// owns all of `Tag::Preference` and writes cold-start users via
|
||
/// `append_legacy_rows`.** This standalone path remains for direct single-vector
|
||
/// embedders; it shares the `encode_legacy_row` writer so the two cannot drift.
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// Returns storage errors from the underlying engine.
|
||
pub fn checkpoint(&self, storage: &dyn crate::storage::StorageEngine) -> crate::Result<()> {
|
||
use crate::{
|
||
schema::EntityId,
|
||
storage::{Tag, WriteBatch, entity_tag_prefix},
|
||
};
|
||
|
||
let prefix = entity_tag_prefix(EntityId::new(0), Tag::Preference);
|
||
let mut batch = WriteBatch::with_capacity(self.inner.len() + 1);
|
||
|
||
// Stage deletion of any stale rows (e.g. a user whose vector was cleared)
|
||
// so the post-swap snapshot holds exactly the current set.
|
||
for item in storage.scan_prefix(&prefix) {
|
||
let (key, _) = item.map_err(crate::schema::TidalError::from)?;
|
||
batch.delete(key);
|
||
}
|
||
|
||
// Write every current row through the single-sourced legacy encoder.
|
||
self.append_legacy_rows(&mut batch, |_| false);
|
||
|
||
storage
|
||
.write_batch(batch)
|
||
.map_err(crate::schema::TidalError::from)?;
|
||
Ok(())
|
||
}
|
||
|
||
/// Restore per-user preference vectors + update counts from a
|
||
/// `Tag::Preference` checkpoint. Skips rows whose stored dimension does not
|
||
/// match this store's `dim` (a schema change) and rows that are torn.
|
||
///
|
||
/// Each restored vector is re-normalized to unit length at the load boundary
|
||
/// (after zeroing any NaN component) so the cosine-scoring invariant holds
|
||
/// even for a torn or tampered row — see the inline comment at the
|
||
/// normalization site.
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// Returns storage errors from the underlying engine.
|
||
pub fn restore(&self, storage: &dyn crate::storage::StorageEngine) -> crate::Result<()> {
|
||
use crate::{
|
||
schema::EntityId,
|
||
storage::{Tag, entity_tag_prefix, parse_key},
|
||
};
|
||
|
||
let prefix = entity_tag_prefix(EntityId::new(0), Tag::Preference);
|
||
let mut count = 0u64;
|
||
for entry in storage.scan_prefix(&prefix) {
|
||
let (key, value) = entry.map_err(crate::schema::TidalError::from)?;
|
||
let Some((_, Tag::Preference, suffix)) = parse_key(&key) else {
|
||
continue;
|
||
};
|
||
if suffix.len() < 8 {
|
||
continue;
|
||
}
|
||
let user_id = u64::from_be_bytes(suffix[0..8].try_into().unwrap_or([0u8; 8]));
|
||
// Decode + NaN-neutralize + re-normalize at the load boundary via the
|
||
// single-sourced decoder (skips torn / dimension-mismatched rows). The
|
||
// re-normalization re-establishes the unit-length invariant that
|
||
// `cosine_similarity` assumes, so a torn row degrades to a zero
|
||
// preference rather than a poisoned non-unit one.
|
||
let Some((update_count, vec)) = decode_legacy_row(&value, self.dim) else {
|
||
continue;
|
||
};
|
||
// Seed the adaptive update count so the learning rate resumes where
|
||
// it was at checkpoint time.
|
||
self.inner.insert(user_id, vec);
|
||
self.update_counts.insert(user_id, update_count);
|
||
let _ = EntityId::new(user_id);
|
||
count += 1;
|
||
}
|
||
if count > 0 {
|
||
tracing::info!(users = count, "preference vectors restored from checkpoint");
|
||
}
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
/// L2-normalize an accumulated preference centroid in place, tolerating a cold
|
||
/// (all-zero) vector.
|
||
///
|
||
/// This is the **one** place the centroid zero-tolerance policy is named. The math
|
||
/// itself lives in [`crate::storage::vector::l2_normalize_in_place`]; this wrapper
|
||
/// only decides what a zero norm *means* here:
|
||
///
|
||
/// A preference centroid is legitimately all-zero until the first signal folds in,
|
||
/// so "no direction yet" is an ordinary state, not an error — the vector is left at
|
||
/// zero and will normalize on the next fold. Embedding writes deliberately do **not**
|
||
/// get this leniency (`storage::vector::lifecycle::ops` propagates the error) because
|
||
/// a zero embedding entering the HNSW index is unrecoverable garbage.
|
||
///
|
||
/// Both `preference` and `multi_preference` call this; before, each carried its own
|
||
/// copy of the arithmetic with a zero threshold ~2900x looser than the canonical one,
|
||
/// which is how the two directions of normalization drifted apart in the first place.
|
||
/// Vectors with `0 < ||v|| < 3.45e-4` are now left alone rather than amplified, since
|
||
/// dividing by a norm that small yields a direction made of rounding error.
|
||
pub(crate) fn normalize_centroid(vec: &mut [f32]) {
|
||
if let Err(err) = l2_normalize_in_place(vec) {
|
||
// A cold or numerically-zero centroid is the ONLY error this can return
|
||
// (see `l2_norm_checked`), and it is an ordinary state here: leave the
|
||
// vector exactly as it was.
|
||
//
|
||
// Anything else is unreachable today. Leaving the centroid untouched stays
|
||
// the safe default if normalization ever grows a second failure mode, but
|
||
// say so out loud rather than discarding it silently.
|
||
if !matches!(err, VectorError::ZeroNormVector) {
|
||
tracing::warn!(error = %err, "unexpected error normalizing preference centroid");
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Canonical legacy single-vector row encoder: `[count:8 LE][dim:4 LE][f32*dim]`.
|
||
///
|
||
/// This is the **one** place the legacy row is written. `append_legacy_rows`, the
|
||
/// standalone [`PreferenceVectors::checkpoint`], and (via `decode_legacy_row`) the
|
||
/// multi-vector cold-start path all go through this pair so the on-disk format
|
||
/// cannot drift across writers/readers.
|
||
#[must_use]
|
||
pub(crate) fn encode_legacy_row(count: u64, vec: &[f32]) -> Vec<u8> {
|
||
let mut value = Vec::with_capacity(8 + 4 + vec.len() * 4);
|
||
value.extend_from_slice(&count.to_le_bytes());
|
||
#[allow(clippy::cast_possible_truncation)]
|
||
value.extend_from_slice(&(vec.len() as u32).to_le_bytes());
|
||
for v in vec {
|
||
value.extend_from_slice(&v.to_le_bytes());
|
||
}
|
||
value
|
||
}
|
||
|
||
/// Canonical legacy single-vector row decoder. Returns `(update_count,
|
||
/// normalized_vector)`, or `None` if the row is torn or its stored dimension does
|
||
/// not match `expected_dim`. Each component is NaN-neutralized and the vector is
|
||
/// re-normalized at the load boundary (the cosine-scoring invariant), mirroring
|
||
/// the restore contract. Paired with [`encode_legacy_row`]; both
|
||
/// [`PreferenceVectors::restore`] and the multi-vector store's legacy decode
|
||
/// delegate here so the format has exactly one encoder and one decoder.
|
||
#[must_use]
|
||
pub(crate) fn decode_legacy_row(value: &[u8], expected_dim: usize) -> Option<(u64, Vec<f32>)> {
|
||
if value.len() < 12 {
|
||
return None;
|
||
}
|
||
let update_count = u64::from_le_bytes(value[0..8].try_into().ok()?);
|
||
let dim = u32::from_le_bytes(value[8..12].try_into().ok()?) as usize;
|
||
if dim != expected_dim || value.len() < 12 + dim * 4 {
|
||
return None;
|
||
}
|
||
let mut vec = Vec::with_capacity(dim);
|
||
for i in 0..dim {
|
||
let off = 12 + i * 4;
|
||
let f = f32::from_le_bytes(value[off..off + 4].try_into().ok()?);
|
||
vec.push(if f.is_nan() { 0.0 } else { f });
|
||
}
|
||
normalize_centroid(&mut vec);
|
||
Some((update_count, vec))
|
||
}
|
||
|
||
#[cfg(test)]
|
||
#[allow(clippy::unwrap_used, clippy::float_cmp)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn set_and_get() {
|
||
let pv = PreferenceVectors::new(3);
|
||
assert!(pv.set(1, vec![3.0, 4.0, 0.0]));
|
||
let v = pv.get(1).unwrap();
|
||
// 3/5, 4/5, 0
|
||
assert!((v[0] - 0.6).abs() < 1e-6);
|
||
assert!((v[1] - 0.8).abs() < 1e-6);
|
||
assert!((v[2] - 0.0).abs() < 1e-6);
|
||
}
|
||
|
||
#[test]
|
||
fn set_wrong_dim_rejected() {
|
||
let pv = PreferenceVectors::new(3);
|
||
assert!(!pv.set(1, vec![1.0, 2.0]));
|
||
assert!(pv.get(1).is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn update_creates_initial() {
|
||
let pv = PreferenceVectors::new(3);
|
||
assert!(pv.update(1, &[1.0, 0.0, 0.0]));
|
||
let v = pv.get(1).unwrap();
|
||
assert!((v[0] - 1.0).abs() < 1e-6);
|
||
}
|
||
|
||
#[test]
|
||
fn update_blends() {
|
||
let pv = PreferenceVectors::with_learning_rate(2, 0.5);
|
||
let _ = pv.set(1, vec![1.0, 0.0]);
|
||
let _ = pv.update(1, &[0.0, 1.0]);
|
||
let v = pv.get(1).unwrap();
|
||
// After blend: (0.5, 0.5), normalized: (1/sqrt(2), 1/sqrt(2))
|
||
let expected = 1.0 / 2.0f32.sqrt();
|
||
assert!((v[0] - expected).abs() < 1e-5);
|
||
assert!((v[1] - expected).abs() < 1e-5);
|
||
}
|
||
|
||
#[test]
|
||
fn cosine_similarity_normalized() {
|
||
let pv = PreferenceVectors::new(3);
|
||
let _ = pv.set(1, vec![1.0, 0.0, 0.0]);
|
||
// Cosine with self = 1.0
|
||
let sim = pv.cosine_similarity(1, &[1.0, 0.0, 0.0]).unwrap();
|
||
assert!((sim - 1.0).abs() < 1e-6);
|
||
// Orthogonal = 0.0
|
||
let sim = pv.cosine_similarity(1, &[0.0, 1.0, 0.0]).unwrap();
|
||
assert!(sim.abs() < 1e-6);
|
||
}
|
||
|
||
#[test]
|
||
fn cosine_similarity_no_pref() {
|
||
let pv = PreferenceVectors::new(3);
|
||
assert!(pv.cosine_similarity(1, &[1.0, 0.0, 0.0]).is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn l2_normalize_zero_vec() {
|
||
let mut v = vec![0.0f32, 0.0, 0.0];
|
||
normalize_centroid(&mut v);
|
||
assert!(v.iter().all(|&x| x == 0.0));
|
||
}
|
||
|
||
/// Pins the ONE behavior change from collapsing the three `l2_normalize`
|
||
/// copies into `storage::vector::l2_normalize_in_place`.
|
||
///
|
||
/// The deleted in-place copies rejected only `||v|| <= f32::EPSILON` (1.2e-7);
|
||
/// the canonical implementation rejects `||v||^2 < f32::EPSILON`, i.e.
|
||
/// `||v|| < 3.45e-4` — about 2900x stricter. A centroid in that gap is now left
|
||
/// alone instead of being scaled to unit length.
|
||
///
|
||
/// That is deliberate: dividing by a norm that small amplifies float noise by
|
||
/// more than 2900x, so the "direction" produced is mostly rounding error.
|
||
/// Leaving the vector as-is is the honest answer, and it matches what embeddings
|
||
/// have always done. Asserted rather than merely documented so the choice cannot
|
||
/// be un-made silently.
|
||
#[test]
|
||
fn normalize_centroid_leaves_numerically_zero_vector_untouched() {
|
||
// ||v|| = 1e-5 * sqrt(3) = 1.73e-5, inside the gap between the two thresholds.
|
||
let mut v = vec![1e-5f32, 1e-5, 1e-5];
|
||
let before = v.clone();
|
||
normalize_centroid(&mut v);
|
||
assert_eq!(
|
||
v, before,
|
||
"a centroid with norm below the canonical threshold must be left as-is, \
|
||
not amplified into a direction made of rounding error"
|
||
);
|
||
|
||
// Just above the threshold it must still normalize, so the guard is a floor
|
||
// and not a silent no-op for real data.
|
||
let mut real = vec![1.0f32, 2.0, 3.0];
|
||
normalize_centroid(&mut real);
|
||
let norm: f32 = real.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||
assert!(
|
||
(1.0 - norm).abs() < 1e-5,
|
||
"real centroid must reach unit norm, got {norm}"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn update_with_custom_rate_blends_at_given_lr() {
|
||
let pv = PreferenceVectors::new(2);
|
||
let _ = pv.set(1, vec![1.0, 0.0]);
|
||
// lr=0.5 means 50% blend.
|
||
let ok = pv.update_with_custom_rate(1, &[0.0, 1.0], 0.5);
|
||
assert!(ok);
|
||
let v = pv.get(1).unwrap();
|
||
let expected = 1.0 / 2.0f32.sqrt();
|
||
assert!((v[0] - expected).abs() < 1e-5);
|
||
assert!((v[1] - expected).abs() < 1e-5);
|
||
}
|
||
|
||
#[test]
|
||
fn update_with_custom_rate_wrong_dim() {
|
||
let pv = PreferenceVectors::new(3);
|
||
assert!(!pv.update_with_custom_rate(1, &[1.0], 0.1));
|
||
}
|
||
|
||
#[test]
|
||
fn len_and_is_empty() {
|
||
let pv = PreferenceVectors::new(3);
|
||
assert!(pv.is_empty());
|
||
assert_eq!(pv.len(), 0);
|
||
let _ = pv.set(1, vec![1.0, 0.0, 0.0]);
|
||
assert!(!pv.is_empty());
|
||
assert_eq!(pv.len(), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn adaptive_lr_decays_monotonically() {
|
||
let pv = PreferenceVectors::new(4);
|
||
let embedding = [1.0f32, 0.0, 0.0, 0.0];
|
||
|
||
// At count=0: alpha = 0.1 / (1 + ln(1)) = 0.1
|
||
let alpha_at_0 = pv.adaptive_lr(1);
|
||
assert!(
|
||
(alpha_at_0 - 0.1).abs() < 1e-5,
|
||
"alpha at count=0 should be 0.1, got {alpha_at_0}"
|
||
);
|
||
|
||
// Apply 1000 updates.
|
||
for _ in 0..1000 {
|
||
let _ = pv.update(1, &embedding);
|
||
}
|
||
|
||
let alpha_at_1000 = pv.adaptive_lr(1);
|
||
|
||
// After 1000 updates, alpha should be substantially smaller.
|
||
assert!(
|
||
alpha_at_1000 < 0.05,
|
||
"alpha at 1000 updates ({alpha_at_1000}) should be < 0.05"
|
||
);
|
||
assert!(
|
||
alpha_at_1000 < alpha_at_0,
|
||
"alpha must decay with update count"
|
||
);
|
||
|
||
// Verify intermediate counts also decay.
|
||
let pv2 = PreferenceVectors::new(4);
|
||
for _ in 0..100 {
|
||
let _ = pv2.update(2, &embedding);
|
||
}
|
||
let alpha_at_100 = pv2.adaptive_lr(2);
|
||
assert!(
|
||
alpha_at_100 < alpha_at_0,
|
||
"alpha at 100 should be < alpha at 0"
|
||
);
|
||
assert!(
|
||
alpha_at_1000 < alpha_at_100,
|
||
"alpha at 1000 should be < alpha at 100"
|
||
);
|
||
}
|
||
|
||
/// Verify the decay formula using actual vector shift magnitudes.
|
||
///
|
||
/// The AC states that after 1000 preference updates, new signals shift the
|
||
/// vector less than at update count=1. This test measures the L2-norm of
|
||
/// `pref_after - pref_before` at different update counts using an orthogonal
|
||
/// interaction direction to maximize the visible shift.
|
||
///
|
||
/// The `ln_1p` decay formula: `alpha = base / (1 + ln(count + 1))` achieves
|
||
/// ~10–15% of the initial shift magnitude at count=1000 (the formula is
|
||
/// logarithmic, not exponential; it cannot achieve < 5% in finite update counts
|
||
/// without an unreasonably small `base_alpha`).
|
||
#[test]
|
||
fn adaptive_lr_shift_magnitude_decays_with_update_count() {
|
||
// Use a fresh instance so we control update counts precisely.
|
||
let pv = PreferenceVectors::new(4);
|
||
|
||
// Helper to measure the L2 shift magnitude of one update in an
|
||
// orthogonal direction from the current preference.
|
||
let measure_shift = |pv: &PreferenceVectors, interaction: &[f32]| -> f32 {
|
||
let before = pv.get(1).unwrap_or_else(|| vec![1.0, 0.0, 0.0, 0.0]);
|
||
let _ = pv.update(1, interaction);
|
||
let after = pv.get(1).unwrap();
|
||
before
|
||
.iter()
|
||
.zip(after.iter())
|
||
.map(|(a, b)| (b - a).powi(2))
|
||
.sum::<f32>()
|
||
.sqrt()
|
||
};
|
||
|
||
// Establish a baseline preference in direction [1, 0, 0, 0].
|
||
let primary = [1.0f32, 0.0, 0.0, 0.0];
|
||
let orthogonal = [0.0f32, 1.0, 0.0, 0.0];
|
||
|
||
// Initial state via set (does not increment update_count).
|
||
assert!(pv.set(1, primary.to_vec()));
|
||
|
||
// Shift at count=0 (first adaptive update).
|
||
let shift_at_0 = measure_shift(&pv, &orthogonal);
|
||
// Reset back to primary to isolate shift measurement.
|
||
assert!(pv.set(1, primary.to_vec()));
|
||
|
||
// Apply 99 more updates to reach count=100.
|
||
for _ in 0..99 {
|
||
let _ = pv.update(1, &primary);
|
||
}
|
||
let shift_at_100 = measure_shift(&pv, &orthogonal);
|
||
assert!(pv.set(1, primary.to_vec()));
|
||
|
||
// Apply another 900 updates to reach count=1000.
|
||
for _ in 0..900 {
|
||
let _ = pv.update(1, &primary);
|
||
}
|
||
let shift_at_1000 = measure_shift(&pv, &orthogonal);
|
||
|
||
// Verify monotonic decay of shift magnitude.
|
||
assert!(
|
||
shift_at_100 < shift_at_0,
|
||
"shift at 100 ({shift_at_100:.6}) must be < shift at 0 ({shift_at_0:.6})"
|
||
);
|
||
assert!(
|
||
shift_at_1000 < shift_at_100,
|
||
"shift at 1000 ({shift_at_1000:.6}) must be < shift at 100 ({shift_at_100:.6})"
|
||
);
|
||
|
||
// Verify the shift at 1000 is substantially smaller than at count=0.
|
||
// The ln_1p formula achieves ~10-15% of initial shift at count=1000.
|
||
let ratio = shift_at_1000 / shift_at_0;
|
||
assert!(
|
||
ratio < 0.20,
|
||
"shift at 1000 ({shift_at_1000:.6}) should be < 20% of shift at 0 ({shift_at_0:.6}); got {ratio:.3}"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn update_count_tracks_calls() {
|
||
let pv = PreferenceVectors::new(2);
|
||
assert_eq!(pv.update_count(1), 0);
|
||
let _ = pv.update(1, &[1.0, 0.0]);
|
||
assert_eq!(pv.update_count(1), 1);
|
||
let _ = pv.update(1, &[0.0, 1.0]);
|
||
assert_eq!(pv.update_count(1), 2);
|
||
// Different user is independent.
|
||
assert_eq!(pv.update_count(2), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn custom_rate_does_not_use_adaptive_lr() {
|
||
let pv = PreferenceVectors::new(2);
|
||
// Set a known state.
|
||
let _ = pv.set(1, vec![1.0, 0.0]);
|
||
// Use custom rate -- should NOT increment update_counts.
|
||
let _ = pv.update_with_custom_rate(1, &[0.0, 1.0], 0.5);
|
||
assert_eq!(
|
||
pv.update_count(1),
|
||
0,
|
||
"update_with_custom_rate should not increment adaptive count"
|
||
);
|
||
}
|
||
|
||
/// Encode a `Tag::Preference` checkpoint value: `[count:8 LE][dim:4 LE][f32*dim]`.
|
||
fn encode_pref_value(count: u64, vec: &[f32]) -> Vec<u8> {
|
||
let mut value = Vec::with_capacity(8 + 4 + vec.len() * 4);
|
||
value.extend_from_slice(&count.to_le_bytes());
|
||
value.extend_from_slice(&u32::try_from(vec.len()).unwrap().to_le_bytes());
|
||
for v in vec {
|
||
value.extend_from_slice(&v.to_le_bytes());
|
||
}
|
||
value
|
||
}
|
||
|
||
/// W12: a torn/tampered checkpoint row that is NOT unit-length must be
|
||
/// re-normalized at the load boundary. `cosine_similarity` assumes the stored
|
||
/// preference is unit-length (it divides only by the candidate norm), so a
|
||
/// non-unit restored vector would mis-scale every similarity for that user.
|
||
/// This covers both failure shapes: a NaN component (zeroed then non-unit) and
|
||
/// a finite-but-non-unit row.
|
||
#[test]
|
||
fn restore_renormalizes_torn_rows_to_unit_length() {
|
||
use crate::{
|
||
schema::EntityId,
|
||
storage::{InMemoryBackend, StorageEngine, Tag, encode_key},
|
||
};
|
||
|
||
let storage = InMemoryBackend::new();
|
||
let dim = 3usize;
|
||
|
||
// User 1: a NaN component. After zeroing, the residual [3,0,4] has norm 5,
|
||
// i.e. NOT unit length — restore must normalize it to [0.6, 0, 0.8].
|
||
let nan_row = encode_pref_value(7, &[3.0, f32::NAN, 4.0]);
|
||
let key1 = encode_key(EntityId::new(0), Tag::Preference, &1u64.to_be_bytes());
|
||
storage.put(&key1, &nan_row).unwrap();
|
||
|
||
// User 2: a finite-but-non-unit row [0, 6, 8] (norm 10). Must normalize to
|
||
// [0, 0.6, 0.8].
|
||
let big_row = encode_pref_value(3, &[0.0, 6.0, 8.0]);
|
||
let key2 = encode_key(EntityId::new(0), Tag::Preference, &2u64.to_be_bytes());
|
||
storage.put(&key2, &big_row).unwrap();
|
||
|
||
// User 3: an all-zero (fully-corrupt) row stays zero (degrades to a zero
|
||
// preference, never a poisoned non-unit one).
|
||
let zero_row = encode_pref_value(1, &[f32::NAN, f32::NAN, f32::NAN]);
|
||
let key3 = encode_key(EntityId::new(0), Tag::Preference, &3u64.to_be_bytes());
|
||
storage.put(&key3, &zero_row).unwrap();
|
||
|
||
let pv = PreferenceVectors::new(dim);
|
||
pv.restore(&storage).unwrap();
|
||
|
||
let v1 = pv.get(1).expect("user 1 restored");
|
||
let n1: f32 = v1.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||
assert!(
|
||
(n1 - 1.0).abs() < 1e-5,
|
||
"user 1 must be unit length, got {n1}"
|
||
);
|
||
assert!((v1[0] - 0.6).abs() < 1e-5 && (v1[2] - 0.8).abs() < 1e-5);
|
||
assert_eq!(pv.update_count(1), 7, "update count must round-trip");
|
||
|
||
let v2 = pv.get(2).expect("user 2 restored");
|
||
let n2: f32 = v2.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||
assert!(
|
||
(n2 - 1.0).abs() < 1e-5,
|
||
"user 2 must be unit length, got {n2}"
|
||
);
|
||
|
||
let v3 = pv.get(3).expect("user 3 restored");
|
||
assert!(
|
||
v3.iter().all(|&x| x == 0.0),
|
||
"a fully-corrupt row degrades to a zero preference, not a poisoned one"
|
||
);
|
||
|
||
// Cosine of the unit-normalized user-1 pref with its own direction is 1.0,
|
||
// proving the invariant cosine_similarity depends on now holds.
|
||
let sim = pv.cosine_similarity(1, &[0.6, 0.0, 0.8]).unwrap();
|
||
assert!(
|
||
(sim - 1.0).abs() < 1e-5,
|
||
"cosine self-similarity must be 1.0, got {sim}"
|
||
);
|
||
}
|
||
|
||
mod proptests {
|
||
use proptest::prelude::*;
|
||
|
||
use super::*;
|
||
|
||
proptest! {
|
||
/// After any sequence of updates, the L2 norm stays approximately 1.0.
|
||
#[test]
|
||
fn l2_norm_invariant(
|
||
updates in proptest::collection::vec(
|
||
proptest::collection::vec(-1.0f32..1.0f32, 4..=4),
|
||
1..20
|
||
),
|
||
) {
|
||
let pv = PreferenceVectors::new(4);
|
||
for emb in &updates {
|
||
let _ = pv.update(1, emb);
|
||
}
|
||
let v = pv.get(1).unwrap();
|
||
let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||
// After a sequence of updates, the vector should be unit-length
|
||
// (within floating-point tolerance) or exactly zero if all inputs
|
||
// collapse to the origin.
|
||
prop_assert!(
|
||
(norm - 1.0).abs() < 1e-4 || norm < f32::EPSILON,
|
||
"norm was {norm}, expected ~1.0"
|
||
);
|
||
}
|
||
}
|
||
}
|
||
}
|