JingsongLi commented on code in PR #82: URL: https://github.com/apache/paimon-vector-index/pull/82#discussion_r3851062894
########## core/src/projected_assign.rs: ########## @@ -0,0 +1,715 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Exact coarse assignment through a low-dimensional projection. +//! +//! `IVFPQIndex::add` must find the nearest coarse centroid of every row. The +//! exact scan is one `n × nlist × d` GEMM, which dominates `add` once PQ +//! encoding is vectorized. This module keeps the result exact but does most of +//! the work in a `d'`-dimensional PCA subspace of the centroids: +//! +//! 1. project the row and centroids with a contractive PCA projection `P` +//! (`d' × d`); +//! 2. for every centroid compute a conservative lower bound on the true +//! squared distance, accounting for the f32 projection and GEMM errors; +//! 3. evaluate exact distances in ascending bound order and stop as soon as +//! the next bound cannot beat the best exact distance (branch-and-bound). +//! +//! Step 2 is a valid lower bound for any contractive `P`, so the assignment +//! never depends on how well the PCA converged; `P` only decides how many +//! exact evaluations step 3 needs. Rows without low-dimensional structure +//! degrade to checking every centroid, i.e. the exact scan plus one small +//! GEMM, so automatic mode only keeps a projection when the centroids are +//! compressible enough (`d' ≤ d / 3` at 95% explained variance). +//! +//! The projection is a training artifact: it is derived deterministically from +//! the centroids, shared by `IVFPQIndex::from_trained`, and never serialized. + +use crate::blas::sgemm_a_bt; +use crate::distance::fvec_l2sqr; +use nalgebra::{DMatrix, SymmetricEigen}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use rayon::prelude::*; + +/// Fraction of centroid variance the projection must explain. +const VARIANCE_TARGET: f64 = 0.95; +/// Automatic mode keeps the projection only when `d' * MAX_AUTO_DP_DIVISOR <= d`. +const MAX_AUTO_DP_DIVISOR: usize = 3; +/// Block subspace iterations; the bound stays valid regardless of convergence. +const SUBSPACE_ITERATIONS: usize = 8; +const SUBSPACE_SEED: u64 = 0x7a5e_c7ed; +/// Smallest projection worth a GEMM; also the rounding unit for `d'`. +const MIN_DP: usize = 8; +/// Rows per parallel block; the projected score matrix is `rows × nlist`. +const MAX_BLOCK_ROWS: usize = 1024; +/// Leave a visible margin below one after rounding the projection to f32. +const CONTRACTION_MARGIN: f64 = 0.999; + +/// How `IVFPQIndex::add` chooses between the exact centroid scan and the +/// projected branch-and-bound. Both produce the exact nearest centroid. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum ProjectedAssignment { + /// Use the projection when the centroids are compressible enough for it + /// to pay off; otherwise scan. + #[default] + Auto, + /// Always build and use the projection (still exact). + Enabled, + /// Always scan. + Disabled, +} + +/// Contractive projection of the coarse centroids plus the per-centroid data +/// the lower bound needs. +#[derive(Debug)] +pub struct CoarseProjection { + d: usize, + dp: usize, + explained_variance: f64, + /// `dp × d`, scaled so its operator norm does not exceed one. + proj: Vec<f32>, + /// `nlist × dp`: projected centroids. + cents_p: Vec<f32>, + /// `|cents_p[c]|²`. + cents_p_norms: Vec<f64>, + /// `|cents_p[c]|`, cached to keep square roots out of assignment's hot loop. + cents_p_norms_sqrt: Vec<f64>, + /// Upper bound on the f32 GEMM error in each projected centroid. + cents_p_errors: Vec<f64>, +} + +impl CoarseProjection { + /// Fit a projection to `nlist` centroids of dimension `d`. + /// + /// Returns `None` when the projection is not worth it (`force == false`) + /// or cannot be built: too few centroids, zero variance, or a `d'` that + /// would not shrink the GEMM by at least `MAX_AUTO_DP_DIVISOR`. + pub(crate) fn train(cents: &[f32], nlist: usize, d: usize, force: bool) -> Option<Self> { + if nlist == 0 || d == 0 { + return None; + } + // Largest d' auto mode may pick, rounded down so the gate below holds; + // forced mode searches up to d / 2 (beyond that the projected GEMM + // costs about as much as the scan it is meant to replace). + let auto_limit = (d / MAX_AUTO_DP_DIVISOR) / MIN_DP * MIN_DP; + let block = if force { + (d / 2).max(MIN_DP).min(d) + } else { + auto_limit + }; + if !force && (block < MIN_DP || nlist < 2 * block) { + return None; + } + let block = block.min(nlist).max(1); + + let mean = column_mean(cents, nlist, d); + let centered: Vec<f32> = cents + .iter() + .enumerate() + .map(|(i, v)| v - mean[i % d]) + .collect(); + let total_variance = centered + .iter() + .map(|v| (*v as f64) * (*v as f64)) + .sum::<f64>() + / nlist as f64; + if total_variance.is_nan() || total_variance <= 0.0 { + return None; + } + + let (basis, eigenvalues) = top_subspace(¢ered, nlist, d, block); + let target = VARIANCE_TARGET * total_variance; + let mut cumulative = 0.0; + let mut dp = None; + for (i, lambda) in eigenvalues.iter().enumerate() { + cumulative += lambda; + if cumulative >= target { + dp = Some(i + 1); + break; + } + } + let dp = match dp { + Some(dp) => dp.div_ceil(MIN_DP) * MIN_DP, + None if force => block, + None => return None, + } + .min(block); + if !force && dp * MAX_AUTO_DP_DIVISOR > d { + return None; + } + let explained_variance = eigenvalues[..dp].iter().sum::<f64>() / total_variance; + let mut proj = basis[..dp * d].to_vec(); + make_contractive(&mut proj, dp, d); + + let mut cents_p = vec![0.0f32; nlist * dp]; + // Project uncentered vectors. Translation cancels in x-c, while this + // avoids a second source of f32 rounding in the lower bound. + sgemm_a_bt(nlist, dp, d, 1.0, cents, &proj, 0.0, &mut cents_p); + let cents_p_norms: Vec<f64> = (0..nlist) + .map(|c| norm_l2sqr_f64(¢s_p[c * dp..(c + 1) * dp])) + .collect(); + let cents_p_norms_sqrt = cents_p_norms.iter().map(|norm| norm.sqrt()).collect(); + let cents_p_errors: Vec<f64> = cents + .chunks_exact(d) + .map(|c| projection_error(norm_upper(c), d, dp)) Review Comment: [P2] Keep Auto translation-invariant Auto decides that this projection will pay off from the centered variance above, but this block projects uncentered centroids and makes `cents_p_errors` proportional to the full `||c||`; line 234 does the same for rows. Adding a common offset leaves every L2 distance unchanged but can inflate those error radii until every lower bound becomes zero. In a release/8-thread repro (`nlist=512,d=96,dp=8,n=4096`, rank-8 centroids transformed as `1e8 + (v-3)*64`), Auto evaluated all 2,097,152 row/centroid pairs and took 7.32 ms versus 1.82 ms for `find_nearest_batch` (4.0x slower). Please make the projection/error calculation safely translation-invariant, or make Auto reject projections based on sampled pruning effectiveness, and cover this translated low-rank case with an evaluation-rate regression. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
