dtourolle/DarkRoom
1
0
Fork 0
Code Issues 39 Pull Requests Actions Packages Projects Releases 20 Wiki Activity

Regroup the library without stopping the window

Browse Source 
Pressing Regroup on a real library did not come back. Clustering 1,813 faces is
the textbook agglomeration — compute every pairwise cosine, then repeatedly scan
all live group pairs, score each with average link, and merge the best — and the
scan is inside the loop. Each merge rescans every surviving pair, and each score
is recomputed from scratch over every cross pair. Some 1.6 million pair scores
per merge, some 700 merges to do.

Three changes, none of which alter the answer.

Only above-threshold pairs can ever matter. An average that reaches the
threshold must have at least one term at or above it, so two groups with no
qualifying pair between them can never merge — not now, and not after any
sequence of merges, since merging only adds terms. The new `neighbours` module
produces exactly that sparse list: 7,875 pairs rather than 1.6 million on the
reference library. It also means the n^2 matrix is never materialised, so memory
goes from O(n^2) to O(edges) — 2.5 GB to a few hundred KB at 25,000 faces.

Merges cannot cross components, so the connected components of that graph are
independent problems: four hundred small agglomerations instead of one large one.

Average link is additive — sum(A u B, C) = sum(A, C) + sum(B, C) — so a merged
group's scores follow by addition. Kept as running (sum, count) per adjacent
pair, a score costs one division instead of a nested loop, and a heap with lazy
invalidation replaces the rescan.

Measured on the reference library: 0.28s, release, for all 1,813 faces.

An exact ANN index was tried and removed, and neighbours.rs records why so it is
not rediscovered as a good idea. IVF with a triangle-inequality bound is exact
and prunes beautifully on synthetic clusters; on real embeddings it prunes
*nothing* — 946 of 946 cell pairs survive. Median pair angle is 88.5 degrees and
the merge threshold is 66.2, so the bound needs cells of radius under ~10
degrees, but two photographs of the same person sit 36-60 degrees apart. No
ball-based partition of a 512-d near-orthogonal space can be tight enough. So
the scan stayed exhaustive and got an unrolled dot product and its blocks spread
across cores instead.

Correctness is held by keeping the old implementation as an oracle: three tests
run both engines over the same population — plain, under co-occurrence and
anchor constraints, and with a size-weighted calibration — and assert the
clusters are identical. Determinism is asserted at a size where the threaded
path is in play.

62 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
dtourolle
2026-08-27 21:16:55 +02:00
co-authored by Claude Opus 5
parent 7dfbe3184a
commit c10dca984f
4 changed files with 1248 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/dr-face/src/neighbours.rs
+566
View File
@@ -0,0 +1,566 @@
//! TRACES: FR-CULL-9 | FR-CULL-10
//! Finding the face pairs that could possibly be the same person.
//!
//! [`crate::cluster`] used to begin by computing every pairwise cosine and
//! holding the lot as an `n²` matrix of `f32`. At 1,800 faces that is 13 MB,
//! which is why it survived; at 25,000 it is 2.5 GB, which is why it could not
//! keep surviving.
//!
//! Almost all of that matrix is thrown away unread. Clustering only ever asks
//! whether a pair is *above* the merge threshold, and in a real library the
//! answer is no for well over 99% of pairs — the reference library's 1,813
//! faces produced 7,875 qualifying pairs out of 1.6 million. So this module
//! answers the only question that is actually asked — **which pairs clear the
//! bar** — and returns that sparse list. Memory goes from `O(n²)` to `O(edges)`
//! and the caller never has to hold a matrix at all.
//!
//! # Exact, not approximate
//!
//! The usual way to make this fast is an approximate nearest-neighbour index,
//! which trades recall for speed: it *misses* some true neighbours, and a
//! missed neighbour here is a face that silently never joins its person.
//! Nothing surfaces that — the screen just quietly shows one person as two —
//! so it is a poor trade for a feature whose whole job is to be trusted. This
//! module is exact, and the clusters it produces are identical to those from a
//! full scan.
//!
//! # An exact index was tried, measured, and removed
//!
//! Worth recording so it is not rediscovered as a good idea. The obvious exact
//! index is IVF with a triangle-inequality bound: group the embeddings into
//! cells, and skip a whole cell **pair** when the geometry proves no member of
//! one can reach any member of the other. On the sphere,
//!
//! ```text
//! angle(x, y) >= angle(c_P, c_Q) - radius(P) - radius(Q)
//! ```
//!
//! so a cell pair is impossible when `cos` of that lower bound falls below the
//! threshold. Exact, no recall loss, and it prunes beautifully on synthetic
//! clusters.
//!
//! It prunes **nothing at all** on real face embeddings. Measured over the
//! 1,813-face reference library, at √n = 43 cells:
//!
//! | quantity | measured |
//! |---|---|
//! | median pair angle | 88.5° (cosine 0.026) |
//! | merge threshold | 66.2° (cosine 0.403) |
//! | median cell radius | 80.4° |
//! | median centroid separation | 85.0° |
//! | cell pairs surviving the bound | **946 of 946 — 100%** |
//!
//! The arithmetic is not close. For the bound to exclude a typical cell pair it
//! needs `radius(P) + radius(Q) < 85° - 66° = 19°`, so cells of radius under
//! ~10°. But two photographs of the *same person* sit 36–60° apart, so even a
//! perfect single-identity cell has a radius three times too large. No
//! ball-based partition of this space can have cells tight enough for the
//! inequality to bite — 512-d embeddings are near-orthogonal, and that is the
//! curse of dimensionality doing exactly what it says.
//!
//! So the scan stayed exhaustive, and the effort went where it actually pays:
//! not materialising the matrix, an unrolled dot product, and spreading the
//! blocks across cores. That is `O(n²)` time and `O(edges)` memory, which for
//! this problem is the honest answer.
use crate::calibrate::Calibration;
/// Rows of the similarity triangle handed to one thread at a time.
///
/// Small enough that the tail of the triangle divides evenly across cores —
/// row `i` does `n - i` comparisons, so equal *row counts* are very unequal
/// work — and large enough that the per-block overhead disappears.
const BLOCK: usize = 64;
/// Below this many faces, do the whole thing on the calling thread.
///
/// Spawning threads for a set this small costs more than the scan.
const THREADS_ABOVE: usize = 2048;
/// One face pair that clears the merge threshold.
///
/// `i < j` always, and the probability is carried because the caller would
/// otherwise recompute the sigmoid it took a dot product to reach.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Pair {
pub i: usize,
pub j: usize,
pub probability: f32,
}
/// What the caller has to tell us about each face.
///
/// Deliberately not [`crate::cluster::Candidate`]: this module has no business
/// knowing what a person or a photograph is, and taking the three arrays it
/// actually reads keeps it testable on bare vectors.
pub struct Faces<'a> {
/// L2-normalised, `EMBEDDING_DIM` long, one per face.
pub embeddings: &'a [Vec<f32>],
/// Source pixels across the aligned crop, for the calibration's size term.
pub crop_px: &'a [f32],
/// Which photograph each face came from. Two faces in one frame are not
/// the same person, so those pairs are never returned (docs/faces.md §9).
pub images: &'a [u64],
}
/// Every pair whose calibrated probability reaches `min_probability`.
///
/// Excludes pairs from the same photograph, which the clusterer would refuse
/// anyway — dropping them here keeps them out of the graph the caller builds
/// and out of the connected components it derives from it.
///
/// Ordered by `(i, j)`, which is what the caller's determinism rests on.
pub fn above_threshold(faces: &Faces, cal: &Calibration, min_probability: f32) -> Vec<Pair> {
let n = faces.embeddings.len();
if n < 2 {
return Vec::new();
}
// The loosest cosine that could clear the bar for *any* pair in the set.
// Cheaper than the sigmoid by far, and it rejects almost everything.
let tau = loosest_cosine(faces.crop_px, cal, min_probability);
let blocks: Vec<(usize, usize)> = (0..n)
.step_by(BLOCK)
.map(|start| (start, (start + BLOCK).min(n)))
.collect();
if n < THREADS_ABOVE {
let mut out = Vec::new();
for &(from, to) in &blocks {
scan_block(faces, cal, min_probability, tau, from, to, &mut out);
}
return out;
}
// One worker per core bar one. This runs on a background thread behind a
// button the user pressed, and NFR-ARCH-2 puts it behind the UI: taking
// every core would stall the window it is reporting progress to.
let workers = std::thread::available_parallelism()
.map(|p| p.get().saturating_sub(1).max(1))
.unwrap_or(1)
.min(blocks.len());
let next = std::sync::atomic::AtomicUsize::new(0);
let mut parts: Vec<Vec<Vec<Pair>>> = std::thread::scope(|scope| {
let handles: Vec<_> = (0..workers)
.map(|_| {
let next = &next;
let blocks = &blocks;
scope.spawn(move || {
// Results stay tagged with their block index, so the order
// of the output does not depend on which thread got there
// first. Determinism is a promise this module keeps.
let mut mine: Vec<(usize, Vec<Pair>)> = Vec::new();
loop {
let b = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let Some(&(from, to)) = blocks.get(b) else {
break;
};
let mut out = Vec::new();
scan_block(faces, cal, min_probability, tau, from, to, &mut out);
mine.push((b, out));
}
mine
})
})
.collect();
let mut slots: Vec<Vec<Vec<Pair>>> = vec![Vec::new(); blocks.len()];
for h in handles {
for (b, pairs) in h.join().unwrap_or_default() {
slots[b].push(pairs);
}
}
slots
});
let mut out = Vec::new();
for slot in &mut parts {
for pairs in slot.drain(..) {
out.extend(pairs);
}
}
out
}
/// Compare rows `from..to` against everything after them.
///
/// The upper triangle, split by rows. Row `i` only looks at `j > i`, so every
/// unordered pair is visited exactly once and the emitted order is `(i, j)`
/// ascending within the block.
fn scan_block(
faces: &Faces,
cal: &Calibration,
min_probability: f32,
tau: f32,
from: usize,
to: usize,
out: &mut Vec<Pair>,
) {
let n = faces.embeddings.len();
for i in from..to {
let a = &faces.embeddings[i];
let crop_a = faces.crop_px[i];
let image_a = faces.images[i];
for j in i + 1..n {
if image_a == faces.images[j] {
continue;
}
let cos = dot(a, &faces.embeddings[j]);
// The cheap rejection, and it takes well over 99% of pairs.
if cos < tau {
continue;
}
let probability = cal.probability(cos, crop_a.min(faces.crop_px[j]), 0.0);
if probability >= min_probability {
out.push(Pair { i, j, probability });
}
}
}
}
/// The lowest cosine that could yield `min_probability` for any pair in the set.
///
/// The calibration is `sigmoid(a·cos + b + w_size·log2(crop))`, so for a fixed
/// size term the cosine boundary is exact. The size term is *not* fixed — it
/// varies per pair with the smaller of the two faces — so the safe bound uses
/// whichever face size pushes the boundary lowest: the largest face when
/// `w_size` is positive, the smallest when it is negative.
///
/// Returns [`f32::NEG_INFINITY`] — a filter that rejects nothing — where no
/// boundary exists: a non-positive steepness, for which probability does not
/// increase with cosine, or a threshold at the ends of the sigmoid. Those are
/// degenerate calibrations rather than impossible ones, and the right response
/// is to stop pruning, not to guess.
fn loosest_cosine(crop_px: &[f32], cal: &Calibration, min_probability: f32) -> f32 {
if cal.a <= 0.0 || !(min_probability > 0.0 && min_probability < 1.0) {
return f32::NEG_INFINITY;
}
let (mut lo, mut hi) = (f32::INFINITY, 0.0_f32);
for &c in crop_px {
let c = c.max(1.0);
lo = lo.min(c);
hi = hi.max(c);
}
if !lo.is_finite() {
return f32::NEG_INFINITY;
}
let extreme = if cal.w_size >= 0.0 { hi } else { lo };
let tau = cal.boundary_at(min_probability, extreme, 0.0);
if tau.is_nan() {
return f32::NEG_INFINITY;
}
// Cosines never exceed 1, so a boundary above it legitimately rejects
// everything. Clamped rather than left free so the comparison stays cheap.
tau.min(1.0)
}
/// Cosine of two L2-normalised embeddings.
///
/// Eight accumulators rather than one. Floating-point addition is not
/// associative, so the compiler may not re-associate a single running total and
/// the loop serialises on the adder's latency; eight independent chains give it
/// something to pipeline and vectorise. The order is fixed and identical on
/// every run, which is what the caller's determinism needs — it is a different
/// order from the naive sum, not a variable one.
pub(crate) fn dot(a: &[f32], b: &[f32]) -> f32 {
const LANES: usize = 8;
let mut acc = [0.0_f32; LANES];
let chunks = a.len() / LANES;
for c in 0..chunks {
let base = c * LANES;
for (l, slot) in acc.iter_mut().enumerate() {
*slot += a[base + l] * b[base + l];
}
}
let mut total =
((acc[0] + acc[1]) + (acc[2] + acc[3])) + ((acc[4] + acc[5]) + (acc[6] + acc[7]));
for k in chunks * LANES..a.len() {
total += a[k] * b[k];
}
total
}
#[cfg(test)]
mod tests {
use super::*;
const DIM: usize = 64;
fn cal() -> Calibration {
Calibration {
a: 30.0,
b: -30.0 * 0.35,
w_size: 0.0,
valid: true,
positive_pairs: 1000,
negative_pairs: 10_000,
}
}
/// A unit vector in a reproducible pseudo-random direction.
///
/// Hashed from the seed rather than drawn from an RNG, so a failure is
/// reproducible from the test alone.
fn vector(seed: u64) -> Vec<f32> {
let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
let mut v = Vec::with_capacity(DIM);
for _ in 0..DIM {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
v.push(((s >> 11) as f64 / (1u64 << 53) as f64) as f32 - 0.5);
}
normalise(v)
}
fn normalise(mut v: Vec<f32>) -> Vec<f32> {
let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
for x in &mut v {
*x /= n;
}
v
}
/// A vector a known cosine away from `base`.
fn near(base: &[f32], other: &[f32], cosine: f32) -> Vec<f32> {
let d = dot(base, other);
let mut perp: Vec<f32> = other.iter().zip(base).map(|(o, b)| o - d * b).collect();
let n = perp.iter().map(|x| x * x).sum::<f32>().sqrt();
for x in &mut perp {
*x /= n;
}
let s = (1.0 - cosine * cosine).max(0.0).sqrt();
normalise(
base.iter()
.zip(&perp)
.map(|(b, p)| cosine * b + s * p)
.collect(),
)
}
struct Set {
embeddings: Vec<Vec<f32>>,
crop_px: Vec<f32>,
images: Vec<u64>,
}
impl Set {
fn faces(&self) -> Faces<'_> {
Faces {
embeddings: &self.embeddings,
crop_px: &self.crop_px,
images: &self.images,
}
}
}
/// `groups` identities, `per` faces each, every face in its own photograph.
fn population(groups: usize, per: usize, tightness: f32) -> Set {
let mut embeddings = Vec::new();
let mut images = Vec::new();
let mut image = 0u64;
for g in 0..groups {
let base = vector(g as u64 + 1);
let off = vector(g as u64 + 9_999);
for m in 0..per {
embeddings.push(if m == 0 {
base.clone()
} else {
near(&base, &off, tightness)
});
images.push(image);
image += 1;
}
}
let crop_px = vec![150.0; embeddings.len()];
Set {
embeddings,
crop_px,
images,
}
}
/// The unpruned, unthreaded, unblocked definition of the answer.
fn reference(faces: &Faces, cal: &Calibration, min_probability: f32) -> Vec<Pair> {
let n = faces.embeddings.len();
let mut out = Vec::new();
for i in 0..n {
for j in i + 1..n {
if faces.images[i] == faces.images[j] {
continue;
}
let cos: f32 = faces.embeddings[i]
.iter()
.zip(&faces.embeddings[j])
.map(|(x, y)| x * y)
.sum();
let p = cal.probability(cos, faces.crop_px[i].min(faces.crop_px[j]), 0.0);
if p >= min_probability {
out.push(Pair {
i,
j,
probability: p,
});
}
}
}
out
}
fn same_pairs(a: &[Pair], b: &[Pair]) -> bool {
a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.i == y.i && x.j == y.j)
}
#[test]
fn nothing_to_pair_is_no_pairs() {
let s = population(1, 1, 0.9);
assert!(above_threshold(&s.faces(), &cal(), 0.9).is_empty());
}
#[test]
fn the_blocked_scan_finds_exactly_what_the_reference_does() {
let s = population(60, 8, 0.97);
let f = s.faces();
let got = above_threshold(&f, &cal(), 0.9);
let want = reference(&f, &cal(), 0.9);
assert!(!want.is_empty(), "the reference found nothing to check");
assert!(same_pairs(&got, &want), "{} vs {}", got.len(), want.len());
}
/// Past `THREADS_ABOVE` the work is split across cores and stitched back
/// together, and the stitching is where an order bug would live.
#[test]
fn the_threaded_scan_finds_exactly_what_the_reference_does() {
let s = population(300, 8, 0.97);
assert!(
s.embeddings.len() > THREADS_ABOVE,
"population is below the threading cutoff"
);
let f = s.faces();
let got = above_threshold(&f, &cal(), 0.9);
let want = reference(&f, &cal(), 0.9);
assert!(!want.is_empty());
assert!(same_pairs(&got, &want), "{} vs {}", got.len(), want.len());
}
/// The size term moves the cosine boundary per pair, so the pre-filter has
/// to be built from the most permissive size in the set or it will drop a
/// pair that would have qualified.
#[test]
fn a_size_weighted_calibration_still_matches_the_reference() {
let mut s = population(60, 8, 0.97);
for (i, c) in s.crop_px.iter_mut().enumerate() {
*c = 40.0 + (i % 17) as f32 * 30.0;
}
let sized = Calibration {
w_size: 0.5,
b: -30.0 * 0.35 - 0.5 * 7.0,
..cal()
};
let f = s.faces();
assert!(same_pairs(
&above_threshold(&f, &sized, 0.9),
&reference(&f, &sized, 0.9)
));
}
/// A negative size weight flips which extreme is permissive. Cheap to get
/// wrong and silent when it is, so it gets its own case.
#[test]
fn a_negative_size_weight_prunes_from_the_other_end() {
let mut s = population(60, 8, 0.97);
for (i, c) in s.crop_px.iter_mut().enumerate() {
*c = 40.0 + (i % 17) as f32 * 30.0;
}
let sized = Calibration {
w_size: -0.5,
b: -30.0 * 0.35 + 0.5 * 7.0,
..cal()
};
let f = s.faces();
assert!(same_pairs(
&above_threshold(&f, &sized, 0.9),
&reference(&f, &sized, 0.9)
));
}
/// A degenerate calibration has no cosine boundary to prune against, and
/// must stop pruning rather than prune on a bound that does not hold.
#[test]
fn a_flat_calibration_prunes_nothing_and_still_agrees() {
let flat = Calibration {
a: 0.0,
b: 4.0,
..cal()
};
assert_eq!(loosest_cosine(&[150.0], &flat, 0.9), f32::NEG_INFINITY);
let s = population(20, 6, 0.97);
let f = s.faces();
assert!(same_pairs(
&above_threshold(&f, &flat, 0.9),
&reference(&f, &flat, 0.9)
));
}
#[test]
fn two_faces_in_one_photograph_are_never_paired() {
let mut s = population(1, 2, 1.0);
s.images = vec![7, 7];
assert!(above_threshold(&s.faces(), &cal(), 0.9).is_empty());
}
#[test]
fn pairs_come_back_in_index_order() {
let s = population(300, 8, 0.97);
let pairs = above_threshold(&s.faces(), &cal(), 0.9);
assert!(pairs
.windows(2)
.all(|w| (w[0].i, w[0].j) < (w[1].i, w[1].j)));
assert!(pairs.iter().all(|p| p.i < p.j));
}
/// Threads must not make the answer depend on which one finished first.
#[test]
fn the_same_input_yields_the_same_pairs() {
let s = population(300, 8, 0.97);
assert_eq!(
above_threshold(&s.faces(), &cal(), 0.9),
above_threshold(&s.faces(), &cal(), 0.9)
);
}
/// The unrolled dot has to agree with the obvious one, tail included — the
/// lengths here are deliberately not multiples of the lane count.
#[test]
fn the_unrolled_dot_matches_the_naive_one() {
for len in [1usize, 7, 8, 9, 63, 64, 65, 512] {
let a: Vec<f32> = (0..len).map(|i| (i as f32 * 0.37).sin()).collect();
let b: Vec<f32> = (0..len).map(|i| (i as f32 * 0.11).cos()).collect();
let naive: f32 = a.iter().zip(&b).map(|(x, y)| x * y).sum();
assert!(
(dot(&a, &b) - naive).abs() < 1e-4,
"len {len}: {} vs {naive}",
dot(&a, &b)
);
}
}
/// The pre-filter is the whole speed story, so it is worth asserting it
/// actually rejects the bulk of the population rather than trusting it to.
#[test]
fn the_threshold_filter_rejects_almost_everything() {
let s = population(60, 8, 0.97);
let n = s.embeddings.len();
let total = n * (n - 1) / 2;
let kept = above_threshold(&s.faces(), &cal(), 0.9).len();
assert!(
kept * 20 < total,
"kept {kept} of {total} pairs, which is not sparse"
);
}
}