dtourolle
8b3abdb787
The embedder's raw output has a length, and the length is a reading of how recognisable the crop was: a blur, an occlusion or a hard profile comes out short. Normalising threw it away. A short vector sits near the middle of the sphere and matches a little of everyone, which is how one bad crop bridges two people in a grouping pass. So the length is kept — the store now holds the raw vector, re-normalised on load, with the length beside it as `faces.quality` — and a face under MIN_GALLERY_QUALITY (14) is a probe: measured against the gallery and placed where it fits, but never what another face is measured against. Two probes are never paired, and a probe is nobody's evidence for a confidence. The People screen shows the number as "Quality 17.3", dimmed below the floor. Faces indexed before this stored unit vectors and have no reading; they are admitted to the gallery, and schema V14 forgets the run marker of every image holding one so the next indexing pass measures them. A peer's unmeasured shard faces are not adopted, or a sync would write that marker back.
342 lines
13 KiB
Rust
342 lines
13 KiB
Rust
//! What an embedder produces, and how it is stored (docs/faces.md §6).
|
||
//!
|
||
//! Deliberately **model-free**: the vector, its identity, its comparison and
|
||
//! its storage encoding are arithmetic, and `calibrate` and `cluster` are built
|
||
//! on them. Keeping them out of the `inference` feature is what lets the part
|
||
//! of this subsystem most likely to be subtly wrong be tested on a machine with
|
||
//! no weights on it.
|
||
//!
|
||
//! [`crate::embed::Embedder`] is the thing that needs a model, and it lives
|
||
//! behind the feature.
|
||
|
||
/// Embedding dimensionality. Fixed by the model family, not a parameter.
|
||
pub const EMBEDDING_DIM: usize = 512;
|
||
|
||
/// The shortest raw embedding a face may be *compared against*.
|
||
///
|
||
/// # What the length of the vector says
|
||
///
|
||
/// ArcFace is trained on the direction of its output and nothing else, and
|
||
/// the length it leaves behind turns out to be a free quality signal: the
|
||
/// magnitude grows with how recognisable the crop was to the model, and a
|
||
/// blurred, occluded, badly lit or hard-profile face comes out short. MagFace
|
||
/// (Meng et al., CVPR 2021) made that the training objective; the plain
|
||
/// ArcFace heads this crate runs already show it, weaker but usable, which is
|
||
/// why it is worth keeping the number the normalisation discards.
|
||
///
|
||
/// # Why it gates the gallery and not the face
|
||
///
|
||
/// A short vector is a bad *reference*: it sits nearer the centre of the
|
||
/// sphere than a real identity does and matches a little of everyone, which
|
||
/// is exactly the face that welds two people together in a clustering pass.
|
||
/// It is not a bad *probe* — the face is still real, still somebody, and
|
||
/// comparing it against good references is the only way it will ever be named.
|
||
/// So a face below this floor is compared against the gallery and never
|
||
/// becomes part of it: see `cluster::Candidate::in_gallery`.
|
||
///
|
||
/// 14 is the operating point for `w600k_mbf`, whose norms on the reference
|
||
/// library run from about 8 on a blur to the high 20s on a clean portrait. A
|
||
/// face whose quality was never recorded — indexed before the number was kept
|
||
/// — is not gated, because a rule that cannot be checked should admit, not
|
||
/// exclude.
|
||
pub const MIN_GALLERY_QUALITY: f32 = 14.0;
|
||
|
||
/// Whether an embedding of this quality may serve as a reference.
|
||
///
|
||
/// `None` is "not measured", and is admitted: the rule is about a number that
|
||
/// was read and found short, not about a number that is missing.
|
||
pub fn in_gallery(quality: Option<f32>) -> bool {
|
||
quality.is_none_or(|q| q >= MIN_GALLERY_QUALITY)
|
||
}
|
||
|
||
/// Which model produced an embedding.
|
||
///
|
||
/// Embeddings from different models are not comparable, and this is the one
|
||
/// mistake that produces plausible-looking garbage rather than an error — so
|
||
/// the id travels *with* the vector rather than beside it, and
|
||
/// [`Embedding::cosine`] refuses a cross-model comparison.
|
||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||
pub struct ModelId(pub std::sync::Arc<str>);
|
||
|
||
impl ModelId {
|
||
pub fn new(s: impl Into<std::sync::Arc<str>>) -> Self {
|
||
Self(s.into())
|
||
}
|
||
pub fn as_str(&self) -> &str {
|
||
&self.0
|
||
}
|
||
}
|
||
|
||
impl std::fmt::Display for ModelId {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
f.write_str(&self.0)
|
||
}
|
||
}
|
||
|
||
/// A 512-d L2-normalised face embedding.
|
||
#[derive(Debug, Clone, PartialEq)]
|
||
pub struct Embedding {
|
||
pub model: ModelId,
|
||
pub v: Box<[f32; EMBEDDING_DIM]>,
|
||
}
|
||
|
||
impl Embedding {
|
||
/// Cosine similarity, which for unit vectors is the plain dot product.
|
||
///
|
||
/// `None` when the two came from different models. That is a real
|
||
/// possibility in a library indexed across a model upgrade, and the
|
||
/// alternative — returning a number — is the failure mode
|
||
/// `faces.model_id` exists to prevent.
|
||
pub fn cosine(&self, other: &Embedding) -> Option<f32> {
|
||
if self.model != other.model {
|
||
return None;
|
||
}
|
||
Some(dot(&self.v, &other.v))
|
||
}
|
||
|
||
/// Storage form: `512 × f16`, 1 KB per face (catalog.md §10.1).
|
||
///
|
||
/// This writes the unit vector. What the catalog stores is the raw one —
|
||
/// `embed::Embedded::to_f16_bytes` — because the length is the quality
|
||
/// and a unit vector has none left to read.
|
||
pub fn to_f16_bytes(&self) -> Vec<u8> {
|
||
self.to_f16_bytes_scaled(1.0)
|
||
}
|
||
|
||
/// The unit vector scaled by `length`, as `512 × f16`.
|
||
pub(crate) fn to_f16_bytes_scaled(&self, length: f32) -> Vec<u8> {
|
||
let mut out = Vec::with_capacity(EMBEDDING_DIM * 2);
|
||
for &x in self.v.iter() {
|
||
out.extend_from_slice(&f32_to_f16_bits(x * length).to_le_bytes());
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Read back from storage, re-normalising.
|
||
///
|
||
/// The f16 round-trip perturbs a unit vector by ~1e-3 in cosine — three
|
||
/// orders below the separation between a match and a non-match — but the
|
||
/// drift is free to remove and invisible if left, so it is removed here
|
||
/// rather than remembered at every call site. The same pass is what turns
|
||
/// a stored raw vector back into the unit one every comparison expects.
|
||
pub fn from_f16_bytes(model: ModelId, bytes: &[u8]) -> Option<Self> {
|
||
read_f16_bytes(model, bytes).map(|(e, _)| e)
|
||
}
|
||
}
|
||
|
||
/// Read a stored vector back, with the length it was stored at.
|
||
///
|
||
/// The length is the quality where the blob is a raw one, and ~1 where it is
|
||
/// a unit vector from before raw vectors were stored — which is why the
|
||
/// catalog keeps the quality beside the blob rather than deriving it from
|
||
/// this: a unit vector reads as a quality of 1, not as "unmeasured".
|
||
pub fn read_f16_bytes(model: ModelId, bytes: &[u8]) -> Option<(Embedding, f32)> {
|
||
if bytes.len() != EMBEDDING_DIM * 2 {
|
||
return None;
|
||
}
|
||
let mut v = Box::new([0.0_f32; EMBEDDING_DIM]);
|
||
for (i, chunk) in bytes.chunks_exact(2).enumerate() {
|
||
v[i] = f16_bits_to_f32(u16::from_le_bytes([chunk[0], chunk[1]]));
|
||
}
|
||
let length = normalise(&mut v);
|
||
Some((Embedding { model, v }, length))
|
||
}
|
||
|
||
fn dot(a: &[f32; EMBEDDING_DIM], b: &[f32; EMBEDDING_DIM]) -> f32 {
|
||
a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
|
||
}
|
||
|
||
/// Scale `v` to unit length, and return the length it had.
|
||
///
|
||
/// The length is the one thing about the raw output that survives being
|
||
/// thrown away by everything downstream, and it is a quality signal
|
||
/// ([`MIN_GALLERY_QUALITY`]) — so it comes back out rather than being lost
|
||
/// here.
|
||
pub(crate) fn normalise(v: &mut [f32; EMBEDDING_DIM]) -> f32 {
|
||
// Clamped rather than checked: a zero-norm embedding is a broken model,
|
||
// not a runtime condition worth an error path, and dividing by 1e-6 keeps
|
||
// the NaN out of the catalog.
|
||
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt().max(1e-6);
|
||
for x in v.iter_mut() {
|
||
*x /= norm;
|
||
}
|
||
norm
|
||
}
|
||
|
||
// ── f16 ───────────────────────────────────────────────────────────────────
|
||
//
|
||
// Hand-rolled rather than pulling in `half`: two functions over a format that
|
||
// has not changed since 2008, used at exactly one boundary. The dependency
|
||
// policy (D13, D1) makes the bar for a new crate high, and this is well under
|
||
// it.
|
||
|
||
fn f32_to_f16_bits(x: f32) -> u16 {
|
||
let bits = x.to_bits();
|
||
let sign = ((bits >> 16) & 0x8000) as u16;
|
||
let exp = ((bits >> 23) & 0xff) as i32 - 127 + 15;
|
||
let mant = bits & 0x007f_ffff;
|
||
|
||
if exp >= 0x1f {
|
||
// Overflow, inf, or NaN. No component of an embedding exceeds its
|
||
// length, and the lengths this model produces are in the tens, so
|
||
// this is the broken-model path; infinity is the honest answer, not a
|
||
// clamp that hides it.
|
||
return sign
|
||
| 0x7c00
|
||
| if mant != 0 && exp == 0x1f + 112 {
|
||
0x200
|
||
} else {
|
||
0
|
||
};
|
||
}
|
||
if exp <= 0 {
|
||
// Subnormal or underflow. A component of a unit 512-vector is ~0.04
|
||
// and a stored one is that times the length, nowhere near here, so
|
||
// this branch exists for correctness rather than for traffic.
|
||
if exp < -10 {
|
||
return sign;
|
||
}
|
||
let mant = mant | 0x0080_0000;
|
||
let shift = (14 - exp) as u32;
|
||
let half = (mant >> shift) as u16;
|
||
// Round to nearest, ties to even.
|
||
let rem = mant & ((1 << shift) - 1);
|
||
let tie = 1 << (shift - 1);
|
||
let round = u16::from(rem > tie || (rem == tie && (half & 1) == 1));
|
||
return sign | (half + round);
|
||
}
|
||
|
||
let half = ((exp as u16) << 10) | (mant >> 13) as u16;
|
||
let rem = mant & 0x1fff;
|
||
let round = u16::from(rem > 0x1000 || (rem == 0x1000 && (half & 1) == 1));
|
||
sign | (half + round)
|
||
}
|
||
|
||
fn f16_bits_to_f32(h: u16) -> f32 {
|
||
let sign = ((h & 0x8000) as u32) << 16;
|
||
let exp = ((h >> 10) & 0x1f) as u32;
|
||
let mant = (h & 0x03ff) as u32;
|
||
|
||
if exp == 0 {
|
||
if mant == 0 {
|
||
return f32::from_bits(sign);
|
||
}
|
||
// Subnormal: renormalise into f32's range.
|
||
let mut e = -1_i32;
|
||
let mut m = mant;
|
||
while m & 0x0400 == 0 {
|
||
m <<= 1;
|
||
e -= 1;
|
||
}
|
||
let m = m & 0x03ff;
|
||
return f32::from_bits(sign | (((127 - 15 + 1 + e) as u32) << 23) | (m << 13));
|
||
}
|
||
if exp == 0x1f {
|
||
return f32::from_bits(sign | 0x7f80_0000 | (mant << 13));
|
||
}
|
||
f32::from_bits(sign | ((exp + 127 - 15) << 23) | (mant << 13))
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
fn unit(seed: u32) -> Embedding {
|
||
let mut v = Box::new([0.0_f32; EMBEDDING_DIM]);
|
||
let mut s = seed.wrapping_mul(2_654_435_761).wrapping_add(1);
|
||
for x in v.iter_mut() {
|
||
s = s.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
|
||
*x = (s >> 8) as f32 / (1u32 << 23) as f32 - 0.5;
|
||
}
|
||
normalise(&mut v);
|
||
Embedding {
|
||
model: ModelId::new("test"),
|
||
v,
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn a_normalised_embedding_has_cosine_one_with_itself() {
|
||
let e = unit(7);
|
||
assert!((e.cosine(&e).unwrap() - 1.0).abs() < 1e-5);
|
||
}
|
||
|
||
#[test]
|
||
fn embeddings_from_different_models_do_not_compare() {
|
||
let a = unit(1);
|
||
let mut b = unit(1);
|
||
b.model = ModelId::new("other");
|
||
assert_eq!(
|
||
a.cosine(&b),
|
||
None,
|
||
"a cross-model cosine must not be a number"
|
||
);
|
||
}
|
||
|
||
/// The claim docs/faces.md §6 makes about the storage format: the f16
|
||
/// round-trip costs ~1e-3 of cosine, three orders below the separation
|
||
/// between a match and a non-match.
|
||
#[test]
|
||
fn f16_round_trip_preserves_the_embedding() {
|
||
for seed in 0..16 {
|
||
let e = unit(seed);
|
||
let back = Embedding::from_f16_bytes(e.model.clone(), &e.to_f16_bytes()).unwrap();
|
||
let cos = e.cosine(&back).unwrap();
|
||
assert!(cos > 0.9999, "seed {seed}: round-trip cosine {cos}");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn normalising_reports_the_length_it_removed() {
|
||
let mut v = Box::new([0.0_f32; EMBEDDING_DIM]);
|
||
v[0] = 3.0;
|
||
v[1] = 4.0;
|
||
let norm = normalise(&mut v);
|
||
assert!((norm - 5.0).abs() < 1e-6, "norm {norm}");
|
||
assert!((v[0] - 0.6).abs() < 1e-6 && (v[1] - 0.8).abs() < 1e-6);
|
||
}
|
||
|
||
/// The gate admits what it cannot measure: a face from before the number
|
||
/// was kept is not a face that was found wanting.
|
||
#[test]
|
||
fn an_unmeasured_quality_is_admitted_to_the_gallery() {
|
||
assert!(in_gallery(None));
|
||
assert!(in_gallery(Some(MIN_GALLERY_QUALITY)));
|
||
assert!(in_gallery(Some(27.5)));
|
||
assert!(!in_gallery(Some(MIN_GALLERY_QUALITY - 0.01)));
|
||
assert!(!in_gallery(Some(8.0)));
|
||
}
|
||
|
||
/// The storage form carries the length, and the length comes back out —
|
||
/// without touching the direction every comparison is made on.
|
||
#[test]
|
||
fn a_raw_vector_round_trips_with_its_length() {
|
||
let e = unit(3);
|
||
let raw = e.to_f16_bytes_scaled(21.5);
|
||
let (back, length) = read_f16_bytes(e.model.clone(), &raw).unwrap();
|
||
assert!((length - 21.5).abs() < 0.05, "length {length}");
|
||
assert!(e.cosine(&back).unwrap() > 0.9999);
|
||
// A unit vector from an older store reads as length 1, not as an
|
||
// error — see `read_f16_bytes` on why that is not "unmeasured".
|
||
let (_, one) = read_f16_bytes(e.model.clone(), &e.to_f16_bytes()).unwrap();
|
||
assert!((one - 1.0).abs() < 1e-2, "length {one}");
|
||
}
|
||
|
||
#[test]
|
||
fn f16_round_trip_rejects_a_wrong_length_blob() {
|
||
assert!(Embedding::from_f16_bytes(ModelId::new("m"), &[0u8; 100]).is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn f16_handles_the_values_an_embedding_actually_contains() {
|
||
// Components of a unit 512-vector cluster around ±1/sqrt(512) ≈ 0.044.
|
||
for &x in &[0.0_f32, 1.0, -1.0, 0.044_194_17, -0.044_194_17, 1e-3, -7e-4] {
|
||
let back = f16_bits_to_f32(f32_to_f16_bits(x));
|
||
assert!(
|
||
(back - x).abs() <= 1e-3 * x.abs().max(1e-3),
|
||
"{x} round-tripped to {back}"
|
||
);
|
||
}
|
||
}
|
||
}
|